#!/usr/bin/env python3
import sys
from collections import defaultdict
import random


def close_scaled(a, b, scale):
    """
    Adjacency test in *scaled* coordinate space (TikZ units we output).
    a, b = (x, y, diam) already scaled by `scale`.
    Two models are adjacent if centers are <= 2 inches apart + their radii.
    """
    # 1 inch = 25.4 mm; in the scaled space, millimeters were multiplied by `scale`
    dist_in = 25.4 * scale
    dx = a[0] - b[0]
    dy = a[1] - b[1]
    # 'a[2]' and 'b[2]' are diameters already scaled
    return dx * dx + dy * dy <= (2 * dist_in + a[2] / 2 + b[2] / 2) ** 2


def components(models):
    """
    Connected components over the *unscaled* millimeter coordinates, using a
    bucket grid of size 211 mm to prune pair checks.
    """
    def mm_close(u, v):
        # Same threshold logic but directly in mm (no 'scale').
        # Allowance matches original: <= (2*254 + 5*(du+dv)) with some scaling slack.
        return ((u[0] - v[0]) ** 2 + (u[1] - v[1]) ** 2) * 100 <= (2 * 254 + 5 * (u[2] + v[2])) ** 2

    GRANULARITY = 211  # mm
    buckets = defaultdict(list)

    for i, (x, y, d) in enumerate(models):
        buckets[(int(x) // GRANULARITY, int(y) // GRANULARITY)].append(i)

    edges = defaultdict(set)
    for (gx, gy), idxs in buckets.items():
        candidates = set()
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                candidates.update(buckets.get((gx + dx, gy + dy), []))
        for a in idxs:
            ua = models[a]
            for b in candidates:
                if b == a:
                    continue
                if mm_close(ua, models[b]):
                    edges[a].add(b)
                    edges[b].add(a)

    comp = {}
    cid = 0
    for s in range(len(models)):
        if s in comp:
            continue
        comp[s] = cid
        q = [s]
        while q:
            nq = []
            for u in q:
                for v in edges[u]:
                    if v in comp:
                        continue
                    comp[v] = cid
                    nq.append(v)
            q = nq
        cid += 1
    return comp


def main():
    if len(sys.argv) < 2:
        print("Usage: draw_board.py INPUT.txt [--grid|--grid211] [--no-models|--no-disks|--no-visuals] [--no-fluff|--minimal]")
        sys.exit(2)

    with open(sys.argv[1], "r") as infile, sys.stdout as texfile:
        # Flags
        draw_grid = any(arg in ("--grid", "--grid211") for arg in sys.argv[2:])
        hide_models = any(arg in ("--no-models", "--no-disks", "--no-visuals") for arg in sys.argv[2:])
        hide_fluff = any(arg in ("--no-fluff", "--minimal") for arg in sys.argv[2:])

        # Read input
        N = int(infile.readline())
        text = None
        if N > 5000:
            text = f"Too many models ({N})"

        models_mm = []
        for _ in range(N):
            # each line: x y diameter   (in millimeters)
            models_mm.append(tuple(map(float, infile.readline().split())))

        # Connected components (unscaled)
        comp = components(models_mm)

        # Colour palette per component index
        colours = {}
        for i in range(len(models_mm)):
            c = comp[i]
            if c == 0:
                colours[i] = "blue"
            elif c == 1:
                colours[i] = "cyan"
            elif c == 2:
                colours[i] = "brown"
            elif c == 3:
                colours[i] = "green"
            elif c == 4:
                colours[i] = "lime"
            elif c == 5:
                colours[i] = "magenta"
            elif c == 6:
                colours[i] = "olive"
            elif c == 7:
                colours[i] = "orange"
            else:
                colours[i] = random.choice(["pink", "purple", "red", "teal", "violet", "yellow"])

        # Compute bounding box (mm) with a small margin of 25.4/2 around disks
        xs = set()
        ys = set()
        for (x, y, d) in models_mm:
            xs.add(x - d / 2 - 25.4 / 2)
            xs.add(x + d / 2 + 25.4 / 2)
            ys.add(y - d / 2 - 25.4 / 2)
            ys.add(y + d / 2 + 25.4 / 2)

        xmin, xmax = min(xs), max(xs)
        ymin, ymax = min(ys), max(ys)
        size_mm = max(xmax - xmin, ymax - ymin)

        # Scale so the figure fits roughly 5x5 cm in TikZ units (1 unit = 1 cm)
        IMAGE_SIZE_CM = 5.0
        scale = IMAGE_SIZE_CM / size_mm  # multiply mm by this to get cm in TikZ

        # Prepare scaled models for drawing
        models = []
        for (x, y, d) in models_mm:
            models.append((x * scale, y * scale, d * scale))

        # Neighbours only needed if we might draw adjacency lines
        if text is None and (not hide_models) and (not hide_fluff):
            neighbours = [[] for _ in range(N)]
            for i, a in enumerate(models):
                neighbours[i] = [b for b in models if a is not b and close_scaled(a, b, scale)]

        # Preamble
        print(
            r"""\documentclass[varwidth]{standalone}
\pdfinfoomitdate=1
\pdftrailerid{}
\pdfsuppressptexinfo=-1
\usepackage{tikz}
\usetikzlibrary{matrix, fit}
\tikzset{every node/.style={font=\sffamily}}
\begin{document}
""",
            file=texfile,
        )

        if text is not None:
            print(text, file=texfile)
            print(r"\end{document}", file=texfile)
            return

        # Begin picture
        print(r"\begin{tikzpicture}", file=texfile)

        # Optional 211x211 mm grid overlay (scaled to TikZ units)
        if draw_grid:
            step_cm = 211 * scale
            x0 = xmin * scale
            x1 = xmax * scale
            y0 = ymin * scale
            y1 = ymax * scale
            xv = x0
            # vertical lines
            while xv <= x1 + 1e-9:
                print(rf"\draw[black!50, thin] ({xv}, {y0}) -- ({xv}, {y1});")
                xv += step_cm
            # horizontal lines
            yh = y0
            while yh <= y1 + 1e-9:
                print(rf"\draw[black!50, thin] ({x0}, {yh}) -- ({x1}, {yh});")
                yh += step_cm
        else:
            # Background light grid(s) for orientation (10 mm and 100 mm in scaled units)
            print(
                rf"\draw[step={scale * 100},gray, thin] ({xmin * scale}, {ymin * scale}) grid ({xmax * scale}, {ymax * scale});",
                file=texfile,
            )
            if scale > 0.04:
                print(
                    rf"\draw[step={scale * 10},gray, very thin] ({xmin * scale}, {ymin * scale}) grid ({xmax * scale}, {ymax * scale});",
                    file=texfile,
                )

        if not hide_models:
            # --- Optional fluff (extra radii + adjacency lines) ---
            if not hide_fluff:
                # 2-inch adjacency aura around each model
                for i, (x, y, diam) in enumerate(models):
                    col = colours[i]
                    print(
                        rf"\draw[fill={col}!10, semitransparent, draw={col}!20] "
                        rf"({x}, {y}) circle ({diam / 2 + 25.4 * scale});"
                    )
                # adjacency lines
                for i, a in enumerate(models):
                    for b in neighbours[i]:
                        print(rf"\draw [blue] ({a[0]}, {a[1]})--({b[0]}, {b[1]});")

            # --- Always draw the disks themselves ---
            for i, (x, y, diam) in enumerate(models):
                # If neighbours are not available (because fluff was hidden), default to 'good'
                if "neighbours" in locals():
                    good = len(neighbours[i]) >= (1 + int(N > 6))
                else:
                    good = True
                colour = "black!20" if good else "red!50"
                print(rf"\draw[fill={colour}, draw] ({x}, {y}) circle ({diam / 2});")

        print(r"\end{tikzpicture}", file=texfile)
        print(r"\end{document}", file=texfile)


if __name__ == "__main__":
    main()