"""Lesson 3 lab: exact linear systems, least squares, and conditioning.

Run:
    python 03-linear-systems-least-squares.py

Requires:
    numpy

The normal-equation function is included to expose the mathematics. Prefer
``numpy.linalg.lstsq`` for real least-squares work because explicitly forming
``A.T @ A`` can worsen numerical conditioning.
"""

from __future__ import annotations

import numpy as np


def solve_two_by_two(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Solve a 2x2 system by elimination, with explicit input checks."""
    if a.shape != (2, 2) or b.shape != (2,):
        raise ValueError(f"Expected A=(2, 2) and b=(2,), received {a.shape} and {b.shape}")

    augmented = np.column_stack((a.astype(float), b.astype(float)))
    if abs(augmented[0, 0]) < abs(augmented[1, 0]):
        augmented[[0, 1]] = augmented[[1, 0]]
    if np.isclose(augmented[0, 0], 0.0):
        raise np.linalg.LinAlgError("Singular system: no usable first pivot")

    augmented[1] -= (augmented[1, 0] / augmented[0, 0]) * augmented[0]
    if np.isclose(augmented[1, 1], 0.0):
        raise np.linalg.LinAlgError("Singular system: no unique solution")

    x = np.empty(2)
    x[1] = augmented[1, 2] / augmented[1, 1]
    x[0] = (augmented[0, 2] - augmented[0, 1] * x[1]) / augmented[0, 0]
    return x


def normal_equation_fit(design: np.ndarray, targets: np.ndarray) -> np.ndarray:
    """Educational least-squares fit via (A.T A)x=A.T b; not for production."""
    if design.ndim != 2 or targets.ndim != 1 or design.shape[0] != targets.shape[0]:
        raise ValueError("Design rows must match the one-dimensional target array")
    return np.linalg.solve(design.T @ design, design.T @ targets)


def main() -> None:
    # Exact system: one intersection of two independent equations.
    exact_a = np.array([[2.0, 1.0], [1.0, -1.0]])
    exact_b = np.array([7.0, 2.0])
    manual = solve_two_by_two(exact_a, exact_b)
    library = np.linalg.solve(exact_a, exact_b)
    assert np.allclose(manual, [3.0, 1.0])
    assert np.allclose(manual, library)
    assert np.allclose(exact_a @ manual, exact_b)

    # Fit y = intercept + slope*x to three observations.
    feature = np.array([0.0, 1.0, 2.0])
    targets = np.array([1.0, 2.0, 2.0])
    design = np.column_stack((np.ones_like(feature), feature))

    normal_fit = normal_equation_fit(design, targets)
    fit, squared_residuals, rank, singular_values = np.linalg.lstsq(
        design, targets, rcond=None
    )
    predictions = design @ fit
    residual = targets - predictions

    assert np.allclose(fit, [7.0 / 6.0, 0.5])
    assert np.allclose(normal_fit, fit)
    assert np.allclose(design.T @ residual, np.zeros(2))
    assert np.allclose(squared_residuals, [np.dot(residual, residual)])
    assert rank == 2

    # Underdetermined: lstsq returns the minimum-norm minimizing solution.
    under_a = np.array([[1.0, 1.0]])
    under_b = np.array([2.0])
    minimum_norm, _, under_rank, _ = np.linalg.lstsq(under_a, under_b, rcond=None)
    pseudoinverse_solution = np.linalg.pinv(under_a) @ under_b
    assert under_rank == 1
    assert np.allclose(minimum_norm, [1.0, 1.0])
    assert np.allclose(pseudoinverse_solution, minimum_norm)
    assert np.allclose(under_a @ minimum_norm, under_b)

    # Every [1-t, 1+t] solves x1+x2=2. The pseudoinverse selects t=0,
    # whose Euclidean norm is no larger than the alternatives.
    alternative = np.array([-2.0, 4.0])
    assert np.allclose(under_a @ alternative, under_b)
    assert np.linalg.norm(minimum_norm) < np.linalg.norm(alternative)

    # Nearly duplicate columns make coefficients sensitive even when predictions fit.
    epsilon = 1e-8
    ill_conditioned = np.array(
        [[1.0, 1.0], [1.0, 1.0 + epsilon], [1.0, 1.0 + 2 * epsilon]]
    )
    condition = np.linalg.cond(ill_conditioned)
    assert condition > 1e8

    print("exact solution:", manual)
    print("least-squares [intercept, slope]:", fit)
    print("predictions:", predictions)
    print("residual:", residual)
    print("A.T @ residual (orthogonality check):", design.T @ residual)
    print("rank and singular values:", rank, singular_values)
    print("minimum-norm underdetermined solution:", minimum_norm)
    print("same solution via pseudoinverse:", pseudoinverse_solution)
    print("ill-conditioned design condition number:", condition)

    # Try it yourself:
    # 1. Add a fourth observation (x=3, y=4) and predict the new fitted line.
    # 2. Center the feature column and compare the condition number.
    # 3. Perturb ill_conditioned targets by 1e-10 and inspect coefficient changes.


if __name__ == "__main__":
    main()