"""Turn theta-rho patterns into G-code for this table.

Machine axes: X = theta in degrees (unlimited, keeps counting), Y = rho in mm (0 = centre).
Patterns never go inside rho_min_mm (3 mm), so they don't press the home switch.
The controller moves X and Y together in a straight line in (theta, rho) space, which on the
sand is an Archimedean spiral piece: exactly how .thr files are meant to be read. Long moves
are still split into short ones so each piece can get its own feed rate: the feed is set so
the *ball* moves at a steady speed, whatever the radius.
"""
import math

DEG = 180.0 / math.pi


def phys_len(t1, r1, t2, r2):
    """Length on the sand (mm) of a theta-rho move, (deg, mm) -> mm. Exact enough for short moves."""
    dt = (t2 - t1) / DEG
    rm = (r1 + r2) / 2
    return math.hypot(rm * dt, r2 - r1)


class Planner:
    """Keeps track of where the ball is and turns patterns into G-code lines."""

    def __init__(self, cfg, theta_deg=0.0, rho_mm=0.0):
        self.cfg = cfg
        self.theta = theta_deg      # current machine X (deg)
        self.rho = rho_mm           # current machine Y (mm)

    # ------------------------------------------------------------ geometry
    def to_machine(self, pts, rotate_deg=0.0, mirror=False):
        """Pattern points (rad, 0..1) -> machine points (deg, mm), not yet shifted by whole turns."""
        m = self.cfg['machine']
        sign = -1.0 if bool(m['invert_theta']) != bool(mirror) else 1.0
        off = m['theta_offset_deg'] + rotate_deg
        R, rmin = m['rho_max_mm'], m['rho_min_mm']
        # rho 0 (the centre) becomes rho_min: under a 9.5 mm ball the 3 mm difference doesn't show
        return [(sign * th * DEG + off, max(rho * R, rmin)) for th, rho in pts]

    def densify(self, mpts):
        mo = self.cfg['motion']
        smax, dmax, smin = mo['max_segment_mm'], mo['max_segment_deg'], mo['min_segment_mm']
        # merge points closer than smin (keeps the first and last point)
        kept = [mpts[0]]
        for p in mpts[1:-1]:
            if phys_len(*kept[-1], *p) >= smin or abs(p[0] - kept[-1][0]) >= dmax:
                kept.append(p)
        if len(mpts) > 1:
            kept.append(mpts[-1])
        out = [kept[0]]
        for (t1, r1), (t2, r2) in zip(kept, kept[1:]):
            # equal steps in (theta, rho) are longest at the outer end, so size them there
            L_outer = math.hypot(max(r1, r2) * (t2 - t1) / DEG, r2 - r1)
            n = max(1, math.ceil(L_outer / smax), math.ceil(abs(t2 - t1) / dmax))
            for i in range(1, n + 1):
                f = i / n
                out.append((t1 + (t2 - t1) * f, r1 + (r2 - r1) * f))
        return out

    def feed(self, t1, r1, t2, r2, speed):
        """G-code F (machine units/min) that gives `speed` mm/s on the sand for this move."""
        D = math.hypot(t2 - t1, r2 - r1)                 # length in machine units (deg, mm mixed)
        L = max(phys_len(t1, r1, t2, r2), 0.02)          # length on the sand
        return min(60.0 * speed * D / L, 200000.0)

    def seconds(self, mpts, speed):
        mo = self.cfg['motion']
        s = 0.0
        for (t1, r1), (t2, r2) in zip(mpts, mpts[1:]):
            s += max(phys_len(t1, r1, t2, r2) / speed, abs(t2 - t1) / mo['max_theta_deg_s'],
                     abs(r2 - r1) / mo['max_rho_mm_s'])
        return s

    # ------------------------------------------------------------ erase
    def erase_path(self, target_theta_start, target_rho, pitch_mm):
        """Spiral from the current ball position to (pattern start), sweeping the whole tray.
        `target_theta_start` may be shifted by any number of whole turns; the shifted value used
        is returned with the path. The spiral turns whichever way unwinds the arm's total angle."""
        R, rmin = self.cfg['machine']['rho_max_mm'], self.cfg['machine']['rho_min_mm']
        t0, r0 = self.theta, self.rho
        far = R if r0 < R / 2 else rmin
        legs = [far] if abs(far - target_rho) > 0.01 * R else []
        legs.append(target_rho)
        if abs(legs[0] - r0) < 0.01 * R and len(legs) > 1:     # already at the far end
            legs = legs[1:]
        direction = -1.0 if t0 > 0 else 1.0
        path = [(t0, r0)]
        t, r = t0, r0
        for i, rr in enumerate(legs):
            turns = abs(rr - r) / pitch_mm
            ideal = t + direction * turns * 360.0
            if i == len(legs) - 1:
                k = round((ideal - target_theta_start) / 360.0)
                tt = target_theta_start + 360.0 * k
            else:
                tt = ideal
            path.append((tt, rr))
            t, r = tt, rr
        return path, path[-1][0] - target_theta_start

    # ------------------------------------------------------------ programs
    def program(self, pts, erase=True, speed=None, rotate_deg=0.0, mirror=False):
        """G-code for one pattern (with the erase spiral before it). Returns
        (lines, segments) where segments = [('erase'|'travel'|'pattern', machine points)]."""
        mo, er = self.cfg['motion'], self.cfg['erase']
        speed = speed or mo['ball_speed_mm_s']
        mp = self.to_machine(pts, rotate_deg, mirror)
        segs = []
        if erase and er['enabled']:
            epath, shift = self.erase_path(mp[0][0], mp[0][1], er['pitch_mm'])
            segs.append(('erase', self.densify(epath), er['speed_mm_s']))
        else:
            shift = 360.0 * round((self.theta - mp[0][0]) / 360.0)
            if phys_len(self.theta, self.rho, mp[0][0] + shift, mp[0][1]) > 0.5:
                segs.append(('travel', self.densify([(self.theta, self.rho), (mp[0][0] + shift, mp[0][1])]), speed))
        mp = [(t + shift, r) for t, r in mp]
        segs.append(('pattern', self.densify(mp), speed))
        lines = []
        for kind, path, v in segs:
            lines.append(f'(--- {kind} ---)')
            for (t1, r1), (t2, r2) in zip(path, path[1:]):
                f = self.feed(t1, r1, t2, r2, v)
                lines.append(f'G1X{t2:.3f}Y{r2:.3f}F{f:.0f}')
        self.theta, self.rho = segs[-1][1][-1]
        return lines, [(k, p) for k, p, _ in segs], sum(self.seconds(p, v) for _, p, v in segs)


HEADER = ['G21', 'G90', 'G94']        # mm, absolute, feed per minute