from dataclasses import dataclass
from typing import List, Optional
import numpy as np
import matplotlib.pyplot as plt

def _prefix_xy(x: np.ndarray, y: np.ndarray):
    def pref(a):
        return np.concatenate([[0.0], np.cumsum(a)])
    return dict(
        Sx=pref(x), Sy=pref(y), Sxx=pref(x * x), Syy=pref(y * y), Sxy=pref(x * y)
    )

def _interval_scatter(pref, i: int, j: int):
    Sx = pref["Sx"][j + 1] - pref["Sx"][i]
    Sy = pref["Sy"][j + 1] - pref["Sy"][i]
    Sxx = pref["Sxx"][j + 1] - pref["Sxx"][i]
    Syy = pref["Syy"][j + 1] - pref["Syy"][i]
    Sxy = pref["Sxy"][j + 1] - pref["Sxy"][i]
    n = j - i + 1
    if n <= 1:
        return 0.0, 0.0, 0.0, n
    Cxx = Sxx - (Sx * Sx) / n
    Cyy = Syy - (Sy * Sy) / n
    Cxy = Sxy - (Sx * Sy) / n
    return Cxx, Cyy, Cxy, n

def _interval_tls_sse(pref, i: int, j: int) -> float:
    Cxx, Cyy, Cxy, n = _interval_scatter(pref, i, j)
    if n <= 1:
        return 0.0
    tr = Cxx + Cyy
    det = Cxx * Cyy - Cxy * Cxy
    disc = max(0.0, tr * tr - 4.0 * det)
    lam_min = 0.5 * (tr - np.sqrt(disc))
    return float(max(0.0, lam_min))

def find_by_dp(
    xy_sorted: np.ndarray, n_breakpoints: int, min_seg_points: int
) -> List[int]:
    N = xy_sorted.shape[0]
    K = max(1, n_breakpoints + 1)
    if N < K * min_seg_points:
        K = max(1, N // max(1, min_seg_points))
    if K <= 1 or N < 2:
        return []
    x = xy_sorted[:, 0]
    y = xy_sorted[:, 1]
    pref = _prefix_xy(x, y)
    dp = np.full((K + 1, N + 1), np.inf, dtype=np.float64)
    prv = np.full((K + 1, N + 1), -1, dtype=np.int32)
    dp[0, 0] = 0.0
    for k in range(1, K + 1):
        j_min = k * min_seg_points
        for j in range(j_min, N + 1):
            i_min = (k - 1) * min_seg_points
            i_max = j - min_seg_points
            best_cost, best_i = np.inf, -1
            for i in range(i_min, i_max + 1):
                if not np.isfinite(dp[k - 1, i]):
                    continue
                sse = _interval_tls_sse(pref, i, j - 1)
                cost = dp[k - 1, i] + sse
                if cost < best_cost:
                    best_cost, best_i = cost, i
            dp[k, j] = best_cost
            prv[k, j] = best_i
    
    cuts = []
    k, j = K, N
    while k > 0:
        i = int(prv[k, j])
        if i < 0:
            break
        cuts.append(i)
        j = i
        k -= 1
    cuts = sorted(cuts)
    interior_starts = [c for c in cuts if 0 < c < N]
    interior_breaks = [s - 1 for s in interior_starts]
    return interior_breaks

# ---------- LineModel ----------
@dataclass
class LineModel:
    # n·x = d (Hessian), ||n||=1
    n: np.ndarray
    d: float
    t: np.ndarray
    c: np.ndarray
    p0: np.ndarray
    p1: np.ndarray

    @staticmethod
    def from_tls(pts: np.ndarray) -> "LineModel":
        c = pts.mean(axis=0)
        A = pts - c # 가장 큰 고유값의 고유벡터가 선의 방향 
        _, vecs = np.linalg.eigh(A.T @ A)
        t = vecs[:, 1] / (np.linalg.norm(vecs[:, 1]) + 1e-12) # 주성분
        n = np.array([-t[1], t[0]]) # t에 수직
        if n.dot(c) < 0: # d>=0 되도록 방향 통일
            n = -n
        d = n.dot(c)
        s = (pts - c) @ t # scalar projections
        smin, smax = np.min(s), np.max(s)
        p0 = c + smin * t
        p1 = c + smax * t
        return LineModel(n=n, d=d, t=t, c=c, p0=p0, p1=p1)
    
    @property
    def length(self) -> float:
        return float(np.linalg.norm(self.p1 - self.p0))
    
    # ---------- piecewise TLS splitter ----------
def piecewise_tls_lines(
    xy: np.ndarray, n_breakpoints: int, min_seg_points: int = 20
):
    xy = np.asarray(xy, dtype=float)
    mask = np.isfinite(xy).all(axis=1)
    xy = xy[mask]
    if xy.shape[0] < max(2, (n_breakpoints + 1) * max(1, min_seg_points)):
        if xy.shape[0] < 2:
            return []
        lm = LineModel.from_tls(xy)
        return [lm]
    interior_breaks = find_by_dp(
        xy, n_breakpoints=n_breakpoints, min_seg_points=min_seg_points
    )
    bounds = [-1] + interior_breaks + [xy.shape[0] - 1]
    lines = []
    pts_list = []
    for k in range(len(bounds) - 1):
        i0 = bounds[k] + 1
        i1 = bounds[k + 1]
        pts = xy[i0 : i1 + 1]
        if pts.shape[0] < 2:
            continue
        lines.append(LineModel.from_tls(pts))
        pts_list.append(pts)
    return lines, interior_breaks

# ---------- minimal demo ----------
def _make_demo_points(n1=120, n2=140, noise=0.01, seed=0):
    rng = np.random.default_rng(seed)
    # seg1: y = 0.2x + 0.5, x in [0, 2.5]
    x1 = np.linspace(0.0, 2.5, n1)
    y1 = 0.2 * x1 + 0.5
    P1 = np.stack([x1, y1], axis=1) + rng.normal(0, noise, size=(n1, 2))
    # seg2: y = -0.1x + 1.2, x in [2.6, 6.5]
    x2 = np.linspace(2.6, 6.5, n2)
    y2 = -0.1 * x2 + 1.2
    P2 = np.stack([x2, y2], axis=1) + rng.normal(0, noise, size=(n2, 2))
    return np.vstack([P1, P2])

if __name__ == "__main__":
    xy = _make_demo_points()
    model, _ = piecewise_tls_lines(xy, n_breakpoints=2, min_seg_points=10)
    print(f"#segments = {len(model)}")
    for i, ln in enumerate(model, 1):
        print(f"seg{i}: len={ln.length:.3f}, t={ln.t}, n={ln.n}, d={ln.d:.3f}")

    plt.figure(figsize=(7, 4.5))
    plt.scatter(xy[:, 0], xy[:, 1], s=10, alpha=0.7, label="points")
    for i, ln in enumerate(model, 1):
        plt.plot([ln.p0[0], ln.p1[0]], [ln.p0[1], ln.p1[1]], label=f"seg{i}")
    plt.title("Piecewise TLS (self-contained)")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.legend()
    plt.tight_layout()
    plt.show()