/*
    Pieter Suurmond, 2006. C++ implementation of Laguerre's root finding method.
    The following Fortran code (http://www.netlib.org/kincaid-cheney/laguerre.f)
    was taken as a starting point:   
    
        c     Numerical Analysis:
        c     The Mathematics of Scientific Computing
        c     D.R. Kincaid & E.W. Cheney
        c     Brooks/Cole Publ., 1990
        c     Section 3.5 Example of Laguerre's Method
        c
              implicit complex (a-h,o-z)
              parameter (n=4, M=7)
              dimension a(0:n)
              data a/(-72.,68.),(1140.,636.),(-216.,190.),(-10.,-26.),(1.,0.)/
              data z/(1.,1.)/
        c
              sn = real(n)
              do 3 k=1,m
                 alpha = a(n)
                 beta = (0.0,0.0)
                 gamma = (0.0,0.0)
        c
                 do 2 j=n-1,0,-1
                    gamma = z*gamma + beta
                    beta  = z*beta  + alpha
                    alpha = z*alpha + a(j)
         2       continue
        c
                 ca = -beta/alpha
                 cb =  ca*ca - 2.0*gamma/alpha
                 c1 = sqrt((sn-1.0)*(sn*cb - ca*ca))
                 cc1 = (ca + c1)/sn
                 cc2 = (ca - c1)/sn
                 cc = cc2
                 if (abs(cc1) .gt. abs(cc2)) then cc = cc1
	            c2 = 1.0/cc
	            z = z + c2
	            d = (ca - cc)/(sn - 1.0)
	            rho1 = sqrt(sn)*c2
	            rho2 = sqrt(sn/(cc*cc + (sn-1)*d*d))
	            print *,'values of z, c2, rho1, and rho2:'
	            print *,z,c2,rho1,rho2
         3    continue
        c

    On UNIX, compile with "g++ -O2 -Wall find_roots.cpp -o find_roots -lcomplex -lm".
    Perhaps use "-lComplex" instead; but always include the math library ("-lm").
*/

#include <iostream>
#include <iomanip>
#include <complex>          // Template-based if we're lucky.

#include "find_roots.hpp"

#define kEPSILON (1e-16)

using namespace std;

/*
    Tries to find one single root.
*/
static int laguerre(int                    n, // Degree = number of coefficients - 1.
                    const complex<double>* a, // Ptr to array of complex coefficients.
                    complex<double>*       z, // Starting/end point for the iteration.
                    int                    m) // Number of iterations.
{
    double  sn(n);    // Notice n = sn = number of coefficients - 1.
    double  snm(n-1);

    for (int k = 1; k <= m; k++)
        {
        complex<double> ca, caca, cb, c1, c2, cc, cc1, cc2;
        complex<double> alpha = a[n];
        complex<double> beta (0.0, 0.0);
        complex<double> gamma(0.0, 0.0);

        for (int j = n - 1; j >= 0; j--)
            {
            gamma = ((*z) * gamma) + beta;
            beta  = ((*z) * beta ) + alpha;
            alpha = ((*z) * alpha) + a[j];
            }
        if (alpha == 0.0)
            {
            cout << "alpha=0!\n";
            return 0;
            }
        ca   = -beta / alpha;
        caca = ca * ca;
        cb   = caca - (2.0 * gamma / alpha);
        c1   = sqrt(snm * ((sn * cb) - caca));
        cc1  = ca + c1;
        cc2  = ca - c1;
        if (abs(cc1) > abs(cc2))
            cc = cc1;
        else
            cc = cc2;
        cc /= sn;
	    c2 = 1.0 / cc;
	    (*z) += c2;             /* Leave result in z. */
	    
        cout << "   z = " << *z  << endl;
        cout << "  c2 = " << c2  << endl;
/*
        complex<double> rho1 = sqrt(sn) * c2;
        complex<double> d    = (ca - cc) / snm;
	    complex<double> rho2 = sqrt(sn / ((cc * cc) + (snm * d * d)));

        cout << "rho1 = " << rho1 << endl;
        cout << "rho2 = " << rho2 << endl;
*/
        }
    return 0;
}


/*
    In most cases, the number of roots will equal the number of input coefficients
    minus one. Only if the last (or latter) coefficient(s) is (are) zero, the number
    of resulting roots may be less.
*/
int find_roots(int                    num_coeffs,   // Degree + 1.
               const complex<double>* coeffs,       // Input array, size=num_coeffs.
               int*                   num_roots,    // Actual degree of the polynome.
               complex<double>*       roots)        // Found roots are written here.
{
    complex<double> *c = NULL, *cc = NULL;
    int             degree = num_coeffs - 1;
    
    *num_roots = 0;
    // Find the ACTUAL degree of the polynome. This guarentees last coeff is non-zero.
    while ((degree > 0) && (coeffs[degree] == 0.0))  
        degree--;
    if (degree < 1)                 // There must at least be 2 coefficients, of which
        return 1;                   // only the first may be zero.
    // Divide-out roots at the origin.
    while ((degree > 0) && (coeffs[0] == 0.0))
        {
        roots[(*num_roots)++] = 0.0;  // Insert root(s) at 0.
        coeffs++;                     // Example:      17       18      17
        degree--;                     //          a17 x  + a18 x   --> x  (a17 + a18 x)
        }                             // Example:               17      17
    // From here on, first and last coeff are non-zero.      a17 x   --> x  (a17)
    if (degree > 0)
        {                                   // Copy input array so we can write in it.
        cc = c = new complex<double>[degree + 1];
        if (!c)
            return 2;                       // Memory failure!
        for (int n = 0; n <= degree; n++)
            c[n] = coeffs[n];
        }
    while (degree > 2)
        {
        int n;                          // See whether all coeffs in between are zero.
        for (n = 1; (n < degree) && (c[n] == 0.0); n++)
            ;                           // Then we have a complex     88
        if (n == degree)                // binominal in this form:   x  = -a0 / a88 
            {                           // Now we have 88 solutions all at once:
            complex<double> p = -c[0] / c[degree];
            double          magnitude = pow(norm(p), 0.5 / static_cast<double>(degree));
            double          angle     = arg(p) / static_cast<double>(degree);
            double          angle_inc = 8.0 * atan(1.0) / static_cast<double>(degree);
            do  {
                roots[(*num_roots)++] = polar(magnitude, angle);
                angle += angle_inc;
                } while (--degree);
            }                           // degree = 0 breaks the while-loop.
        else
            {                           // Solve numerically.
            complex<double> z(-64.0, -64.0);
            
            int e = laguerre(degree,    // Number of coefficients - 1.
                             c,         // Ptr to array of complex coefficients.
                             &z,        // Found root and starting point for the iteration.
                             20);       // Number of iterations.
            if (e)
                { delete[] cc; return e; }
            // Synthetic division.
            degree--;                           // c[degree] stays the same.
	        for (int i = degree; i >= 0; i--)
                c[i] += c[i+1] * z;
            if (abs(c[0]) >= kEPSILON)          // Must be (almost) zero.
                { delete[] cc; return 4; }
            c++;                                // Inc. pointer to skip first coeff.
            roots[(*num_roots)++] = z;          
            }
        }
    if (degree == 2)                // Solve quadratic equation analytically:
        {                                               //        --------
        complex<double> a2 = c[2] * 2.0;                //       / 2
        complex<double> D  = sqrt(c[1] * c[1] -         // D = \/ b - 4ac
                                  4.0 * c[2] * c[0]);   //
        roots[(*num_roots)++] = (-c[1] + D) / a2;       // x1 = (-b+D)/2a
        roots[(*num_roots)++] = (-c[1] - D) / a2;       // x1 = (-b-D)/2a
        }
    else if (degree == 1)                     // Solve linear equation analytically:
        {                                     // a0 + a1 x = 0  <=>  x = -a0 / a1.
        roots[(*num_roots)++] = -c[0] / c[1]; // Division by zero can never occur.
        }
    // else (if (degree == 0)) we're done.
    if (c)
        delete[] cc;                          // Release temporary workspace.
    return 0;
}


int main()
{
/*                                                          3               
    For example, a 3th-degree polynome: (x+1)(x+2)(x-3) = 1x - 7x - 6
*/
    const int       n_coeffs = 4;       // INPUT
    complex<double> a[n_coeffs];

    a[0] = complex<double>(-6.0, 0.0);  // Supply coeffs in ascending degree.
    a[1] = complex<double>(-7.0, 0.0);
    a[2] = complex<double>( 0.0, 0.0);
    a[3] = complex<double>( 1.0, 0.0);

    int             n_roots;            // OUTPUT
    complex<double> r[n_coeffs - 1];    // Allocate one less than input.

    int e = find_roots(n_coeffs,        // Degree + 1.
                       a,               // Input array, size=n_coeffs.

                       &n_roots,        // Receive actual degree of polynome.
                       r);
    if (e)
        cout << "ERROR: find_roots()=" << e << ".\n";
    else
        {
        cout << "Roots:\n" << setiosflags(ios::fixed | ios::right) << setprecision(18);
        for (int n = 0; n < n_roots; n++)
            cout << "r" << n << " = " << setw(40) << r[n] << "\n";
        }
    return e;
}