(require hyrule [-> case unless])
(import hyrule [dec inc])
(import curses math random sys)

(setv WIDTH None)
(setv SHAPE-CHARS (list "#$%&*+/=?"))
(setv COL-DEPTH 15)

(defn gx [x] (if (isinstance x tuple) (get x 0) (lfor k x (gx k))))

(defn gy [x] (if (isinstance x tuple) (get x 1) (lfor k x (gy k))))

(defn in-bounds [x y w h] (and (>= x 0) (>= y 0) (< x w) (< y h)))

(defn rotate [shape qta]
    (sfor #(x y) shape #(
        (+ (* (** -1 (// qta 2)) (% (inc qta) 2) x) (* -1 (** -1 (// qta 2)) (% qta 2) y))
        (+ (* (** -1 (// qta 2)) (% (inc qta) 2) y) (* (** -1 (// qta 2)) (% qta 2) x)))))

(defn translate [shape dx dy]
    (sfor #(x y) shape #((+ x dx) (+ y dy))))

(defn normalise [shape]
    (translate shape (- (min (gx shape) :default 0)) (- (min (gy shape) :default 0))))

(defn init-grid [w h]
    (dfor x1 (range w) y1 (range h)
        #(x1 y1) (sfor #(x2 y2) [#((dec x1) y1) #(x1 (dec y1)) #((inc x1) y1) #(x1 (inc y1))]
            :if (in-bounds x2 y2 w h) #(x2 y2))))

(defn draw [scr shape ch [bg " "] [w 0] [h 0]]
    (if (isinstance scr curses.window)
        (if shape (for [#(x y) shape] (scr.addch y x ch)) "")
        (if shape (let [cg (or scr (lfor _ (range (or h (inc (max (gy shape)))))
                (lfor _ (range (or w (inc (max (gx shape))))) bg)))]
            (for [#(x y) shape] (setv (get cg y x) ch))
            cg) scr)))

(defn to-pbm [shapes]
    (setv w (+ 2 (max (lfor s shapes (- (max (gx s)) (min (gx s)) -1)))))
    (setv h (inc (sum (lfor s shapes (inc (- (max (gy s)) (min (gy s)) -1))))))
    (setv header f"P1\n{w} {h}\n")
    (+ header (.join "\n" (lfor s shapes
        (.join "\n" (lfor r (draw "" (translate (normalise s) 1 1) "1" "0" w) (.join "" r)))))
        "\n" (* "0" w) "\n"))

(defn colour [t]
    (defn ud [x] (let [y (-> x (% 0.5) (* 6) (min 1))]
        (round (* COL-DEPTH (if (>= (% x 1) 0.5) (- 1 y) y)))))
    (let [r (ud (+ t (/ 1 3))) g (ud t) b (ud (+ t (/ 2 3)))] f"{r} {g} {b}"))

(defn to-ppm [shapes]
    (setv w (+ 3 (max (lfor s shapes (max (gx s))))))
    (setv h (+ 3 (max (lfor s shapes (max (gy s))))))
    (setv header f"P3\n{w} {h}\n{COL-DEPTH}\n")
    (setv grid (draw "" #{#(0 0)} "0 0 0" "0 0 0" w h))
    (for [#(i s) (enumerate shapes)]
        (setv grid (draw grid (translate s 1 1) (colour (/ i (len shapes))))))
    (+ header (.join "\n" (lfor r grid (.join " " r))) "\n"))

(defn nbhd [xy]
    (lfor d [#(1 0) #(0 1) #(-1 0) #(0 -1)] #((+ (gx xy) (gx d)) (+ (gy xy) (gy d)))))

(defn hop-shape [xy src dst]
    (setv (get dst xy) #{})
    (del (get src xy))
    (for [sc (nbhd xy)]
        (when (in sc dst) (.add (get dst xy) sc) (.add (get dst sc) xy))
        (when (in sc src) (.discard (get src sc) xy))))

(defn components [shape]
    (setv pl (set (shape.keys)))
    (setv comps [])
    (setv comp [])
    (setv fi 0)
    (while pl
        (if (>= fi (len comp))
            (do (when comp (comps.append (set comp))) (setv comp [(pl.pop)]) (setv fi 0))
            (do
                (for [ap (nbhd (get comp fi))]
                    (when (in ap pl) (comp.append ap) (pl.remove ap)))
                (+= fi 1))))
    comps)

(defn equiv [s1 s2]
    (let [ns2 (normalise s2)]
        (any (lfor r (range 4) (not (.symmetric-difference (normalise (rotate s1 r)) ns2))))))

(defn builder [stdscr]
    (setv KND {
        (ord "h") #(-1 0) (ord "t") #(0 1)
        (ord "n") #(0 -1) (ord "s") #(1 0)
        curses.KEY-LEFT #(-1 0) curses.KEY-DOWN #(0 1)
        curses.KEY-UP #(0 -1) curses.KEY-RIGHT #(1 0)})
    (stdscr.clear)
    (setv epi 0)
    (setv cursor #(0 0))
    (setv k 0)
    (setv message "")
    (while (!= k (ord "q"))
        (stdscr.refresh)
        (when (in k KND)
            (let [d (get KND k) w #((+ (gx cursor) (gx d)) (+ (gy cursor) (gy d)))]
                (when (in-bounds (gx w) (gy w) WIDTH WIDTH) (setv cursor w))))
        (case (chr k)
            "r" (do
                (when (in cursor (get shapes 0))
                    (unless epi (setv epi (len shapes)) (shapes.append {})))
                (setv rpc 0)
                (while (and
                    (< (len (get shapes epi)) WIDTH)
                    (setx opts (list (sfor rxy (| (.keys (get shapes epi)) #{cursor}) nxy (nbhd rxy)
                        :if (in nxy (get shapes 0)) nxy))))
                    (setv cursor (random.choice opts))
                    (hop-shape cursor (get shapes 0) (get shapes epi))
                    (+= rpc 1))
                (setv message f"randomly picked {rpc}/{WIDTH} tiles ({(len (get shapes epi))}, {opts})"))
            "a" (unless epi (when (in cursor (get shapes 0))
                (setv epi (len shapes))
                (shapes.append {}))
                (setv message f"started shape {epi}")))
        (when epi
            (when (in cursor (get shapes 0)) (hop-shape cursor (get shapes 0) (get shapes epi)))
            (when (>= (len (get shapes epi)) WIDTH)
                (setv message f"completed shape {epi}")
                (when (.startswith (setx message (cond
                    (let [MPW (-> WIDTH (math.sqrt) (* 2) (math.floor)) ns (normalise (get shapes epi))]
                        (any (lfor p ns (>= (max (gx p) (gy p)) MPW))))
                        "WRN: shape would be too elongated"
                    (any (lfor s (cut shapes None -1) (equiv s (get shapes epi))))
                        "WRN: shape would be a duplicate"
                    (any (lfor c (components (get shapes 0)) (% (len c) WIDTH)))
                        "ERR: shape would leave invalid-size parts"
                    True message)) "ERR")
                    (for [xy (list (get shapes epi))]
                        (hop-shape xy (get shapes epi) (get shapes 0)))
                    (shapes.pop))
                (setv epi 0)
                (for [xy (get shapes 0)]
                    (when (= (len (get shapes 0 xy)) 1)
                        (setv epi (len shapes))
                        (shapes.append {})
                        (setv opts [xy])
                        (while (and
                            (= (len opts) 1)
                            (< (len (get shapes epi)) WIDTH))
                            (setv cursor (get opts 0))
                            (setx opts (list (get shapes 0 cursor)))
                            (hop-shape cursor (get shapes 0) (get shapes epi)))
                        (unless (message.startswith "ERR")
                            (setv message f"inferred {(len (get shapes epi))} cells of forced shape {epi}"))
                        (when (any (lfor comp (components (get shapes 0))
                            (- (len (get shapes epi)) (% (len comp) WIDTH))))
                            (setv message f"ERR: invalid-size parts forced?"))
                        (break)))
                (unless (or epi (not (get shapes 0)) (in cursor (get shapes 0)))
                    (setv cursor (next (iter (lfor x (range WIDTH) y (range WIDTH)
                        :if (in #(x y) (get shapes 0)) #(x y))))))))
        (when (= (len (get shapes 0)) 0) (break))
        (for [#(i s) (enumerate shapes)]
            (draw stdscr s (get SHAPE-CHARS i)))
        (stdscr.addch (gy cursor) (gx cursor) "@")
        (stdscr.hline WIDTH 0 " " (get (stdscr.getmaxyx) 1))
        (stdscr.addstr WIDTH 0
            f"({(gx cursor)}, {(gy cursor)}) in poly {(try (get (lfor #(i s) (enumerate shapes) :if (in cursor s) i) 0) (except [IndexError] "?"))}{(if message (+ " | " message) "")}")
        (setv k (stdscr.getch))))

(defn load-shapes [fname]
    (setv fl (with [fh (open fname "r")]
        (lfor l (fh.readlines) :if (not (l.startswith "#")) (l.strip))))
    (setv WIDTH (len (get fl 0)))
    (setv shapes [(init-grid WIDTH WIDTH)])
    (for [#(y r) (enumerate fl) #(x c) (enumerate r) :if (.isdigit c)]
        (let [i (int c)]
            (when (<= (len shapes) i) (shapes.extend (lfor _ (range (- (inc i) (len shapes))) {})))
            (hop-shape #(x y) (get shapes 0) (get shapes i))))
    shapes)
 
(try
    (setv mode (get sys.argv 1))
    (setv fname (get sys.argv 2))
    (cond
        (.isdigit mode) (do
            (setv WIDTH (int mode))
            (setv shapes [(init-grid WIDTH WIDTH)])
            (curses.wrapper builder)
            (setv og (draw "" #{#(0 0)} "_" "_" WIDTH WIDTH))
            (for [#(i s) (enumerate shapes)] (setv og (draw og s (str i) "_")))
            (with [fh (open fname "w")] (fh.write (+ (.join "\n" (lfor r og (.join "" r))) "\n"))))
        (= mode "q") (print (to-pbm (lfor s (load-shapes fname)
            :if s (rotate s (random.randrange 4)))))
        (= mode "a") (print (to-ppm (lfor s (load-shapes fname) :if s s))))
    (except [FileNotFoundError]
        (print f"{(get sys.argv 0)}: can't load from {fname}: No such file or directory")))