"""
Maths Crossword Generator
=========================
Usage:
    python3 crossword_builder.py EIGHT THREE NINE SEVEN TEN FOUR
    python3 crossword_builder.py --clues "EIGHT:4+4,THREE:9-6,..." NINE SEVEN TEN FOUR
    python3 crossword_builder.py --output my_puzzle.pdf EIGHT THREE NINE SEVEN TEN FOUR

The script will:
  1. Find a valid crossword layout for the given words
  2. Generate a printable PDF
"""

import sys
import random
import argparse
from copy import deepcopy


# ============================================================
# STEP 1: CROSSWORD SOLVER
# ============================================================

def find_intersections(w1, w2):
    """Return list of (i, j, char) where w1[i] == w2[j]."""
    return [(i, j, c) for i, c in enumerate(w1) for j, c2 in enumerate(w2) if c == c2]


def build_grid(placements):
    """Build a dict {(row,col): char} from placements."""
    grid = {}
    for word, is_across, r, c in placements:
        for i, ch in enumerate(word):
            pos = (r, c + i) if is_across else (r + i, c)
            grid[pos] = ch
    return grid


def conflicts(placements, word, is_across, r, c):
    """Return True if placing word here conflicts with existing placements."""
    grid = build_grid(placements)
    word_len = len(word)

    # Cell immediately before start must be empty
    before = (r, c - 1) if is_across else (r - 1, c)
    if before in grid:
        return True

    # Cell immediately after end must be empty
    after = (r, c + word_len) if is_across else (r + word_len, c)
    if after in grid:
        return True

    # Track which direction(s) occupy each cell
    cell_dirs = {}
    for ew, eia, er, ec in placements:
        d = 'across' if eia else 'down'
        for k in range(len(ew)):
            pos = (er, ec + k) if eia else (er + k, ec)
            cell_dirs.setdefault(pos, set()).add(d)

    new_dir = 'across' if is_across else 'down'

    for i, ch in enumerate(word):
        pos = (r, c + i) if is_across else (r + i, c)

        if pos in grid:
            if grid[pos] != ch:
                return True  # character mismatch
            if new_dir in cell_dirs.get(pos, set()):
                return True  # same-direction word already here
        else:
            # Empty cell: perpendicular neighbors must also be empty
            if is_across:
                n1, n2 = (r - 1, c + i), (r + 1, c + i)
            else:
                n1, n2 = (r + i, c - 1), (r + i, c + 1)
            if n1 in grid or n2 in grid:
                return True

    return False


def is_connected(placements):
    """Check that all placed words form a single connected component (BFS)."""
    if len(placements) <= 1:
        return True
    n = len(placements)
    adj = [set() for _ in range(n)]
    for i in range(n):
        w1, ia1, r1, c1 = placements[i]
        pos1 = {(r1, c1+k) if ia1 else (r1+k, c1) for k in range(len(w1))}
        for j in range(i + 1, n):
            w2, ia2, r2, c2 = placements[j]
            pos2 = {(r2, c2+k) if ia2 else (r2+k, c2) for k in range(len(w2))}
            if pos1 & pos2:
                adj[i].add(j)
                adj[j].add(i)
    visited = {0}
    queue = [0]
    while queue:
        curr = queue.pop()
        for nb in adj[curr]:
            if nb not in visited:
                visited.add(nb)
                queue.append(nb)
    return len(visited) == n


def count_intersections(placements):
    """Count total number of cell-sharing intersections between words."""
    total = 0
    for i, (w1, ia1, r1, c1) in enumerate(placements):
        pos1 = {(r1, c1 + k) if ia1 else (r1 + k, c1) for k in range(len(w1))}
        for w2, ia2, r2, c2 in placements[i+1:]:
            pos2 = {(r2, c2 + k) if ia2 else (r2 + k, c2) for k in range(len(w2))}
            total += len(pos1 & pos2)
    return total


def grid_bounds(placements):
    """Return (min_r, min_c, max_r, max_c) of all placements."""
    cells = []
    for word, is_across, r, c in placements:
        for i in range(len(word)):
            cells.append((r, c + i) if is_across else (r + i, c))
    rows = [p[0] for p in cells]
    cols = [p[1] for p in cells]
    return min(rows), min(cols), max(rows), max(cols)


def normalize(placements):
    """Shift placements so grid starts at (0,0)."""
    if not placements:
        return placements
    min_r, min_c, _, _ = grid_bounds(placements)
    return [(w, ia, r - min_r, c - min_c) for w, ia, r, c in placements]


def solve(words, max_attempts=5000):
    """
    Try to find a valid crossword layout for the given words.
    Returns list of (word, is_across, row, col) or None.
    """
    words = [w.upper() for w in words]
    best = None
    best_score = (-1, -1)  # (words_placed, intersections)

    for attempt in range(max_attempts):
        random.shuffle(words)
        placements = []

        # Place first word horizontally at origin
        first = words[0]
        placements.append((first, True, 0, 0))

        for word in words[1:]:
            placed = False
            # Try to intersect with each already-placed word
            candidates = []
            for existing_word, is_across, er, ec in placements:
                for ei, ej, ch in find_intersections(existing_word, word):
                    # existing_word[ei] == word[ej] == ch
                    # Place new word perpendicular to existing
                    new_across = not is_across
                    if is_across:
                        # existing is horizontal, new word is vertical
                        # existing cell at (er, ec+ei), new word col = ec+ei, row = er - ej
                        nr, nc = er - ej, ec + ei
                    else:
                        # existing is vertical, new word is horizontal
                        # existing cell at (er+ei, ec), new word row = er+ei, col = ec - ej
                        nr, nc = er + ei, ec - ej

                    if not conflicts(placements, word, new_across, nr, nc):
                        candidates.append((word, new_across, nr, nc))

            if candidates:
                choice = random.choice(candidates)
                placements.append(choice)
                placed = True

            if not placed:
                # Try placing with adjacency search (wider search)
                break

        placements = normalize(placements)
        score = (len(placements), count_intersections(placements))

        if score > best_score:
            best_score = score
            best = placements

        # Perfect solution: all words placed and connected
        if len(placements) == len(words) and is_connected(placements):
            return normalize(placements)

    return best


# ============================================================
# STEP 2: PDF GENERATOR
# ============================================================

def generate_pdf(placements, clues_map, output_path):
    try:
        from reportlab.lib.pagesizes import A4
        from reportlab.lib import colors
        from reportlab.pdfgen import canvas as rl_canvas
        from reportlab.lib.units import cm
    except ImportError:
        print("ERROR: reportlab not installed. Run: pip install reportlab")
        sys.exit(1)

    DARK_BLUE = colors.Color(0.18, 0.18, 0.55)
    RED       = colors.Color(0.85, 0.07, 0.07)

    placements = normalize(placements)
    min_r, min_c, max_r, max_c = grid_bounds(placements)
    rows = max_r + 1
    cols = max_c + 1

    # Build letter grid
    grid = [[None]*cols for _ in range(rows)]
    for word, is_across, r, c in placements:
        for i, ch in enumerate(word):
            if is_across:
                grid[r][c+i] = ch
            else:
                grid[r+i][c] = ch

    # Number clue starts
    starts = {}
    for word, is_across, r, c in placements:
        starts.setdefault((r, c), [])
    clue_nums = {}
    n = 1
    for r in range(rows):
        for c in range(cols):
            if (r, c) in starts:
                clue_nums[(r, c)] = str(n)
                n += 1

    # Sort words into across/down with their numbers
    across_words = sorted([(w, r, c) for w, ia, r, c in placements if ia],     key=lambda x: (x[1], x[2]))
    down_words   = sorted([(w, r, c) for w, ia, r, c in placements if not ia], key=lambda x: (x[1], x[2]))

    page_w, page_h = A4
    margin = 2.2 * cm
    cv = rl_canvas.Canvas(output_path, pagesize=A4)

    # Border
    cv.setStrokeColor(DARK_BLUE)
    cv.setLineWidth(3)
    cv.rect(1.2*cm, 1.2*cm, page_w - 2.4*cm, page_h - 2.4*cm, fill=0, stroke=1)

    # Title
    cv.setFont("Helvetica-Bold", 28)
    cv.setFillColor(RED)
    cv.drawString(margin, page_h - 2.8*cm, "Crossword")

    # Subtitle
    cv.setFont("Helvetica-Bold", 14)
    cv.setFillColor(DARK_BLUE)
    cv.drawString(margin, page_h - 4.2*cm, "Write the answers to this puzzle in words:")
    cv.drawString(margin, page_h - 5.0*cm, "ONE, TWO, THREE, ...")

    # Grid sizing
    max_grid_w = page_w - 2 * margin - 1*cm
    max_grid_h = page_h * 0.42
    cell_size = min(int(max_grid_w / cols), int(max_grid_h / rows), 38)
    grid_w = cols * cell_size
    grid_x = (page_w - grid_w) / 2
    grid_top_y = page_h - 6.0*cm

    # Draw grid
    for r in range(rows):
        for col in range(cols):
            cx = grid_x + col * cell_size
            cy = grid_top_y - r * cell_size
            if grid[r][col] is not None:
                cv.setFillColor(colors.white)
                cv.setStrokeColor(DARK_BLUE)
                cv.setLineWidth(1.2)
                cv.rect(cx, cy - cell_size, cell_size, cell_size, fill=1, stroke=1)
                if (r, col) in clue_nums:
                    cv.setFillColor(DARK_BLUE)
                    cv.setFont("Helvetica-Bold", 10)
                    cv.drawString(cx + 2, cy - 11, clue_nums[(r, col)])

    # Outer grid border
    cv.setStrokeColor(DARK_BLUE)
    cv.setLineWidth(2.5)
    cv.rect(grid_x, grid_top_y - rows * cell_size, grid_w, rows * cell_size, fill=0, stroke=1)

    grid_bottom_y = grid_top_y - rows * cell_size

    # Clues
    clue_y = grid_bottom_y - 0.7*cm
    col1_x = margin
    col2_x = page_w / 2 + 0.3*cm

    cv.setFont("Helvetica-Bold", 15)
    cv.setFillColor(DARK_BLUE)
    cv.drawString(col1_x, clue_y, "Across")
    cv.drawString(col2_x, clue_y, "Down")
    clue_y -= 0.55*cm

    line_h = 0.62*cm

    for i, (w, r, c) in enumerate(across_words):
        num = clue_nums.get((r, c), "?")
        eq = clues_map.get(w, "_ + _ = ?")
        y = clue_y - i * line_h
        cv.setFont("Helvetica-Bold", 13); cv.setFillColor(RED)
        cv.drawString(col1_x, y, f"{num}.")
        cv.setFillColor(DARK_BLUE)
        cv.drawString(col1_x + 0.9*cm, y, eq)

    for i, (w, r, c) in enumerate(down_words):
        num = clue_nums.get((r, c), "?")
        eq = clues_map.get(w, "_ + _ = ?")
        y = clue_y - i * line_h
        cv.setFont("Helvetica-Bold", 13); cv.setFillColor(RED)
        cv.drawString(col2_x, y, f"{num}.")
        cv.setFillColor(DARK_BLUE)
        cv.drawString(col2_x + 0.9*cm, y, eq)

    cv.save()
    print(f"PDF saved: {output_path}")


# ============================================================
# STEP 3: MAIN
# ============================================================

NUMBER_WORDS = {
    'ONE': 1, 'TWO': 2, 'THREE': 3, 'FOUR': 4, 'FIVE': 5,
    'SIX': 6, 'SEVEN': 7, 'EIGHT': 8, 'NINE': 9, 'TEN': 10,
}


def generate_equation(word):
    """Generate a +/- equation with 2 or 3 operations for a number word (for ages ~6)."""
    n = NUMBER_WORDS.get(word.upper())
    if n is None:
        return "_ + _ = ?"

    num_ops = random.choice([2, 3])

    for _ in range(500):
        # Build operators ensuring at least one + and one -
        ops = [random.choice(['+', '-']) for _ in range(num_ops)]
        if len(set(ops)) < 2:
            ops[random.randrange(num_ops)] = '-' if ops[0] == '+' else '+'

        # Pick first num_ops numbers randomly (1–9), derive the last
        nums = [random.randint(1, 9) for _ in range(num_ops)]
        total = nums[0]
        valid = True
        for k in range(1, num_ops):
            total = total + nums[k] if ops[k - 1] == '+' else total - nums[k]
            if total < 0:
                valid = False
                break
        if not valid:
            continue

        # Last number must satisfy: total OP last = n
        last = (n - total) if ops[-1] == '+' else (total - n)
        if 1 <= last <= 9:
            nums.append(last)
            parts = [str(nums[0])]
            for op, num in zip(ops, nums[1:]):
                parts.append(f"{op} {num}")
            return " ".join(parts) + " = ?"

    # Fallback to simple addition
    for a in range(1, 10):
        b = n - a
        if 1 <= b <= 9:
            return f"{a} + {b} = ?"
    return f"= {n}"


def parse_clues(clues_str):
    """Parse 'EIGHT:4+4,THREE:9-6' into {'EIGHT': '4+4', 'THREE': '9-6'}"""
    result = {}
    for part in clues_str.split(","):
        if ":" in part:
            word, eq = part.split(":", 1)
            result[word.strip().upper()] = eq.strip()
    return result


def main():
    parser = argparse.ArgumentParser(description="Generate a maths crossword PDF from a list of words.")
    parser.add_argument("words", nargs="+", help="Words to place in the crossword (e.g. EIGHT THREE NINE)")
    parser.add_argument("--clues", default="", help="Equations per word: 'EIGHT:4+4,THREE:9-6,...'")
    parser.add_argument("--output", default="crossword.pdf", help="Output PDF filename")
    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducible layouts")
    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)

    # If a single integer is given, pick that many random number words
    if len(args.words) == 1 and args.words[0].isdigit():
        count = int(args.words[0])
        all_words = list(NUMBER_WORDS.keys())
        if count > len(all_words):
            print(f"ERROR: Only {len(all_words)} number words available (ONE–TEN).")
            sys.exit(1)
        words = random.sample(all_words, count)
    else:
        words = [w.upper() for w in args.words]
    clues_map = parse_clues(args.clues)

    # Fill in auto-generated clues for any word without one
    for w in words:
        if w not in clues_map:
            clues_map[w] = generate_equation(w)

    print(f"Solving layout for: {', '.join(words)}")
    placements = solve(words)

    if not placements:
        print("ERROR: Could not find a valid layout. Try different words or fewer words.")
        sys.exit(1)

    placed_words = [w for w, _, _, _ in placements]
    missing = [w for w in words if w not in placed_words]
    if missing:
        print(f"WARNING: Could not place: {', '.join(missing)}")

    print(f"Placed {len(placements)}/{len(words)} words")
    print(f"Layout:")
    for word, is_across, r, c in sorted(placements, key=lambda x: (x[2], x[3])):
        direction = "across" if is_across else "down"
        print(f"  {word:12} {direction:6}  row={r}, col={c}")

    generate_pdf(placements, clues_map, args.output)


if __name__ == "__main__":
    main()