' ******************************************************************
' Utility functions to handle roots of polynomials
' ******************************************************************

' ------------------------------------------------------------------
' Type definition
' ------------------------------------------------------------------

TYPE Complex
  X AS DOUBLE
  Y AS DOUBLE
END TYPE

' ******************************************************************

FUNCTION SetRealRoots(Deg AS INTEGER,_
                      Z() AS Complex,_ 
                      Tol AS DOUBLE) AS INTEGER
' ------------------------------------------------------------------
' Set the imaginary part of a root to zero if it is less than a
' fraction Tol of its real part. This root is therefore considered
' real. The function returns the total number of real roots.
' ------------------------------------------------------------------

  DIM AS INTEGER I, N

  FOR I = 1 TO Deg
    IF Z(I).Y <> 0 AND ABS(Z(I).Y) < Tol * ABS(Z(I).X) THEN 
      Z(I).Y = 0
    END IF
  NEXT I

  ' Count real roots
  N = 0
  FOR I = 1 TO Deg
    IF Z(I).Y = 0 THEN N = N + 1
  NEXT I

  SetRealRoots = N
END FUNCTION

SUB SortRoots(Deg AS INTEGER, Z() AS Complex)
' ------------------------------------------------------------------
' Sorts roots so that:
'
' (1) The N real roots are stored in elements 1..N of vector Z,
'     in increasing order.
'
' (2) The complex roots are stored in elements (N + 1)..Deg of
'     vector Z and are unordered.
' ------------------------------------------------------------------

  DIM AS DOUBLE A

  DIM AS INTEGER I, J, K, Nr, Nc
   
  ' Count real and complex roots
  Nr = 0: Nc = 0
  FOR I = 1 TO Deg
    IF Z(I).Y = 0 THEN Nr = Nr + 1 ELSE Nc = Nc + 1
  NEXT I

  IF Nr > 0 THEN REDIM R(1 TO Nr) AS DOUBLE
  IF Nc > 0 THEN REDIM C(1 TO Nc) AS Complex

  ' Store real roots in R and complex roots in (X,Y)
  J = 0: K = 0
  FOR I = 1 TO Deg
    IF Z(I).Y = 0 THEN
      J = J + 1
      R(J) = Z(I).X
    ELSE
      K = K + 1
      C(K).X = Z(I).X
      C(K).Y = Z(I).Y
    END IF
  NEXT I

  ' Sort real roots
  FOR I = 1 TO Nr - 1
    K = I: A = R(I)
    FOR J = I + 1 TO Nr
      IF R(J) < A THEN K = J: A = R(J)
    NEXT J
    SWAP R(I), R(K)
  NEXT I  
       
  ' Transfer real roots into elements 1..Nr
  FOR I = 1 TO Nr
    Z(I).X = R(I)
    Z(I).Y = 0
  NEXT I

  ' Transfer complex roots into elements (Nr+1)..Deg
  FOR I = 1 TO Nc
    J = I + Nr
    Z(J).X = C(I).X
    Z(J).Y = C(I).Y
  NEXT I
END SUB
