#!/usr/bin/env python3
import argparse
import numpy as np
import matplotlib.pyplot as plt


knoise = 0.1


def make_data(n=60, seed=0):
    rng = np.random.default_rng(seed)

    x = np.linspace(0, 1, n)
    y_true = 1.0 + 2.0 * x - 1.5 * x**2
    y = y_true + knoise * rng.normal(size=n)

    # 外れ値を追加
    outlier_idx = [10, 25, 45]
    y[outlier_idx] += np.array([0.8, -1.0, 0.9])

    return x, y, y_true, outlier_idx


def design_matrix(x, degree):
    return np.vstack([x**k for k in range(degree + 1)]).T


def predict(x, coef):
    return design_matrix(x, len(coef) - 1) @ coef


def least_squares_fit(x, y, degree):
    X = design_matrix(x, degree)
    coef, *_ = np.linalg.lstsq(X, y, rcond=None)
    return coef


def weighted_least_squares_fit(x, y, degree, weights):
    X = design_matrix(x, degree)
    Wsqrt = np.sqrt(weights)

    Xw = X * Wsqrt[:, None]
    yw = y * Wsqrt

    coef, *_ = np.linalg.lstsq(Xw, yw, rcond=None)
    return coef


def huber_weights(residual, delta):
    a = np.abs(residual)

    weights = np.ones_like(a)
    mask = a > delta

    # Huber IRLS:
    # 小さい残差: w = 1
    # 大きい残差: w = delta / |r|
    weights[mask] = delta / a[mask]

    return weights


def huber_irls_fit(x, y, degree, delta=0.15, maxiter=50, tol=1.0e-10):
    coef = least_squares_fit(x, y, degree)

    history = []

    for it in range(maxiter):
        residual = predict(x, coef) - y
        weights = huber_weights(residual, delta)

        coef_new = weighted_least_squares_fit(x, y, degree, weights)

        diff = np.linalg.norm(coef_new - coef)

        history.append({
            "iteration": it + 1,
            "coef": coef_new.copy(),
            "max_abs_residual": float(np.max(np.abs(residual))),
            "min_weight": float(np.min(weights)),
            "coef_change": float(diff),
        })

        coef = coef_new

        if diff < tol:
            break

    return coef, weights, history


def main():
    parser = argparse.ArgumentParser(
        description="Huber regression demo using IRLS"
    )
    parser.add_argument("--degree", type=int, default=2)
    parser.add_argument("--delta", type=float, default=0.15)
    parser.add_argument("--maxiter", type=int, default=50)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    x, y, y_true, outlier_idx = make_data(seed=args.seed)

    coef_l2 = least_squares_fit(x, y, args.degree)
    coef_huber, weights, history = huber_irls_fit(
        x,
        y,
        args.degree,
        delta=args.delta,
        maxiter=args.maxiter,
    )

    print("[least squares]")
    print("coef =", coef_l2)
    print("max abs residual =", np.max(np.abs(predict(x, coef_l2) - y)))

    print()
    print("[Huber IRLS]")
    print("coef =", coef_huber)
    print("max abs residual =", np.max(np.abs(predict(x, coef_huber) - y)))
    print("iterations =", len(history))

    print()
    print("[final weights]")
    for i, w in enumerate(weights):
        if w < 0.999:
            print(f"index={i:2d}, x={x[i]:.3f}, y={y[i]:.3f}, weight={w:.4f}")

    xx = np.linspace(0, 1, 400)

    fig, (ax1, ax2, ax3) = plt.subplots(
        3,
        1,
        figsize=(9, 9),
        sharex=True,
        gridspec_kw={"height_ratios": [3, 1.4, 1.2]},
    )

    ax1.scatter(x, y, c="black", s=28, label="data")
    ax1.scatter(
        x[outlier_idx],
        y[outlier_idx],
        facecolors="none",
        edgecolors="red",
        s=120,
        linewidths=2,
        label="injected outliers",
    )

    ax1.plot(xx, 1.0 + 2.0 * xx - 1.5 * xx**2, "k:", label="true curve")
    ax1.plot(xx, predict(xx, coef_l2), color="tab:orange", lw=2, label="least squares")
    ax1.plot(xx, predict(xx, coef_huber), color="tab:blue", lw=2, label="Huber IRLS")

    ax1.set_ylabel("y")
    ax1.set_title(
        f"Huber regression by IRLS, degree={args.degree}, delta={args.delta}"
    )
    ax1.grid(True, alpha=0.3)
    ax1.legend()

    r_l2 = predict(x, coef_l2) - y
    r_huber = predict(x, coef_huber) - y

    ax2.axhline(0, color="gray", lw=1)
    ax2.plot(x, r_l2, "o-", color="tab:orange", label="least squares residual")
    ax2.plot(x, r_huber, "o-", color="tab:blue", label="Huber residual")
    ax2.axhline(args.delta, color="gray", ls="--", lw=1)
    ax2.axhline(-args.delta, color="gray", ls="--", lw=1)
    ax2.set_ylabel("residual")
    ax2.grid(True, alpha=0.3)
    ax2.legend()

    ax3.plot(x, weights, "o-", color="tab:green")
    ax3.set_xlabel("x")
    ax3.set_ylabel("Huber weight")
    ax3.set_ylim(-0.05, 1.05)
    ax3.grid(True, alpha=0.3)

    fig.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()