# Implementation of module {affine}.
# Last edited on 2021-09-24 09:23:47 by stolfi

import affine
import trafo
import rn
import rmxn
import sys
from math import sqrt, sin, cos, floor, ceil, inf, nan, pi

class Affine_IMP(trafo.Trafo):

  def __init__(self,name,d,Adir,bdir,Ainv,binv):
    trafo.Trafo.__init__(self,name)
    self.d = d
    self.A = (Adir, Ainv)
    self.b = (bdir, binv)
    # ....................................................................

  def combine(self,T2):
    if T2 == None:
      return self
    else:
      assert type(T2) is tuple and len(T2) == 2
      assert t2[0] == +1 or T2[0] == -1
      if isinstance(T2[1], affine.Affine):
        return affine.compose(self, T2)
      else
        return (+1, self, T2)
    # ....................................................................

  # ----------------------------------------------------------------------


def make(name,d,Adir,bdir,Ainv,binv):
  if name == None:
    name = "??"
  else:
    assert type(name) is str
  assert is_square(Adir, d), "{Adir} has wrong size of shape";
  assert is_square(Ainv, d), "{Ainv} has wrong size of shape";
  assert is_long(bdir, d), "{bdir} has wrong size"
  assert is_long(binv, d), "{binv} has wrong size"
  return affine.Affine(name,d,Adir,bdir,Ainv,binv)
  # ----------------------------------------------------------------------

def dimension(A):
  kinv, Aobj = unpack(A)
  return Aobj.d
  # ----------------------------------------------------------------------

def apply(p,A):
  kinv, Aobj = unpack(A)
  j = (1-kinv)//2;
  assert len(p) == Aobj.d
  if Aobj.A[j] != None: p = rmxn.map_row(p,Aobj.A[j])
  if Aobj.b[j] != None: p = rn.add(p, Aobj.b[j])
  return p
  # ----------------------------------------------------------------------

def unpack(A):
  if isinstance(A, affine.Affine):
    return +1, A
  else:
    assert type(A) is tuple, "{A} is neither {Affine} obj nor tuple"
    assert len(A) == 2, "length of {A} tuple is not 2"
    kinv, Aobj = A
    assert Aobj != None
    assert type(kinv) is int
    assert kinv == +1 or kinv == -1
    assert isinstance(Aobj, affine.Affine), "{A[1]} is not {Affine} obj"
    return kinv, Aobj
  # ----------------------------------------------------------------------

def lin_matrix(A):
  kinv, Aobj = unpack(A)
  j = (1-kinv)//2;
  return Aobj.A[j]
  # ----------------------------------------------------------------------

def disp_vector(A):
  kinv, Aobj = unpack(A)
  j = (1-kinv)//2;
  return Aobj.b[j]
  # ----------------------------------------------------------------------

def inv(A):
  kinv, Aobj = unpack(A)
  if kinv == -1:
    return Aobj
  else:
    return (-1, Aobj)
  # ----------------------------------------------------------------------

def compose(A1,A2):
  d = dimension(A1);
  assert dimension(A2) == d, "incompatible dimensions"
  M = [None,None]
  b = [None,None]
  for j in range(2):
    if j == 0:
      B1 = A1; B2 = A2
    else:
      B1 = inv(A2); B2 = inv(A1)
    M1 = lin_matrix(B1)
    M2 = lin_matrix(B2)
    if M1 == None:
      M[j] = M2
    elif M2 == None:
      M[j] = M1 
    else:
      M[j] = rmxn.mul(M1, M2)
    b1 = disp_vector(B1);
    b2 = disp_vector(B2);
    if b1 == None:
      b[j] = b2
    else:
      if M2 != None:
        b1 = rmxn.map_row(b1,M2)
      if b2 == None:
        b[j] = b1
      else:
        b[j] = rn.add(b1, b2)
  return make(None, d, M[0], b[0], M[1], b[1]) 
  # ----------------------------------------------------------------------

# AUXILIARY PROCS

def is_square(M,d):
  # Returns true iff the matrix {M} is {None} or has {d} rows and {d} columns.
  if M == None: return True
  if len(M) != d: return False
  for k in range(d):
    if len(M[k]) != d: return False
  return True
  # ----------------------------------------------------------------------

def is_long(v,d):
  # Returns true iff the vector {v} is {None} or has {d} elements.
  if v == None: return True
  if len(v) != d: return False
  return True
  # ----------------------------------------------------------------------