from dataclasses import dataclass

import numpy as np

from geometry.vector import Vector


@dataclass
class AffineTransform:
  A: np.ndarray
  A_extended: np.ndarray

  inv_A: np.ndarray
  inv_A_extended: np.ndarray

  def __init__(self, A: np.ndarray, b: Vector):
    n = len(A)

    assert A.shape == (n, n)
    assert b.a.shape == (n, 1)

    self.A = np.array(A, dtype=float)
    self.A.setflags(write=False)

    self.A_extended = np.zeros((n + 1, n + 1), dtype=float)
    self.A_extended[:-1, :-1] = self.A
    self.A_extended[:-1, -1:] = b.a
    self.A_extended[-1, -1] = 1
    self.A_extended.setflags(write=False)

    self.inv_A = np.linalg.inv(self.A)
    self.inv_A_extended = np.linalg.inv(self.A_extended)

  def apply_shift(self, v: Vector) -> Vector:
    v_extended = np.vstack((v.a, np.eye(1)))
    result_extended = self.A_extended @ v_extended
    return Vector(result_extended[:-1, :])

  def apply_no_shift(self, v: Vector) -> Vector:
    return Vector(self.A @ v.a)

  def invert_shift(self, v: Vector) -> Vector:
    v_extended = np.vstack((v.a, np.eye(1)))
    result_extended = self.inv_A_extended @ v_extended
    return Vector(result_extended[:-1, :])

  def invert_no_shift(self, v: Vector) -> Vector:
    return Vector(self.inv_A @ v.a)
