/*
    PZT_coeff.c
    PoZeTools v 0.52, march 31, 2005. Copyright (c) 2002-2005 Pieter Suurmond

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to
    deal in the Software without restriction, including without limitation the
    rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
    sell copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions: 

      The above copyright notice and this permission notice shall be included in
      all copies or substantial portions of the Software. 
      Any person wishing to distribute modifications to the Software is 
      requested to send the modifications to the original developer so that they
      can be incorporated into the canonical version. 

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND ON INFRINGEMENT. 
    IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
    DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR 
    OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE 
    USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

#include <stdlib.h>
#include <stdio.h>              /* For FILE. */
#include <string.h>             /* For strcmp(). */
#include <math.h>

#include "PZT_complex_math.h"
#include "PZT_complex_text.h"
#include "PZT_poze.h"           /* PZT_coeff.h needs POZE. */
#include "PZT_coeff.h"

/*------------------------------------------------------------------------------*/
/*
    Allocates array of coefficients that must be released after use by the caller.
    Called by COEFF_calc() twice.
*/
static int multiply_roots(cartesian_array* roots,  /* Array of roots (poles or zeros). */
                          coeff_array*     coeffs, /* Receives here! */
                          FILE*            msg)    /* May be NULL. */
{
    int*    idx;
    int     i, j, a;
    
    coeffs->coeff = NULL;
    idx = malloc((1 + roots->number) * sizeof(int));  /* One extra to prevent allocating 0 bytes, not really used though! */
    if (!idx)
        return 1;                           /* Not enough temporary memory! */
    coeffs->number = roots->number + 1;     /* The maximum number of coefficients, some may appear to null later. */
    coeffs->coeff  = malloc(coeffs->number * sizeof(cartesian_coeff));
    if (!coeffs->coeff)
        { free(idx); return 2; }            /* Not enough memory for coeffs! */
    coeffs->number = 0;                     /* Count only the non-zero coefficients. */
    for (i = 0; i <= roots->number; i++)
        {
        coeffs->coeff[coeffs->number].cartesian.im = 0.0;
        if (msg) fprintf(msg, "  coeff%d =", i);
        if (i)
            {
            coeffs->coeff[coeffs->number].cartesian.re = 0.0; /* Imaginary part already cleared. */
            a = 0; 
            idx[a] = 0;
            do  {
                complexCartesian    product;
                for (j = a + 1; j < i; j++)   /* Init following indices. */
                    idx[j] = 1 + idx[j - 1];
                product.im = 0.0;
                if (i & 1)
                    {
                    product.re = -1.0;
                    if (msg) fprintf(msg, " - ");
                    }
                else
                    {
                    product.re = 1.0;
                    if (msg) fprintf(msg, " + ");
                    }
                for (j = 0; j < i; j++)
                    {
                    cartesianMultiply(&roots->cartesian[idx[j]], &product);
                    if (msg)
                        {
                        if (j) fprintf(msg, "*");
                        fprintf(msg, "R%d", idx[j]);
                        }
                    }                                                    /* Accumulate. */
                cartesianAdd(&product, &coeffs->coeff[coeffs->number].cartesian);
                a = i-1;                                     /* Select rightmost index. */
                while ((a>=0) && (++idx[a] + (i-a) > roots->number))     /* Inc. index. */
                    a--;                         /* Select preceding index on overflow. */
                } while (a >= 0);
            }
        else
            {        /* Imaginary part is already cleared. */
            coeffs->coeff[coeffs->number].cartesian.re = 1.0;
            if (msg) fprintf(msg, " 1");
            }
        /* Automatically delete coefficients that are null. */
        if (cartesianMagnitude(&coeffs->coeff[coeffs->number].cartesian) >= kPZT_epsilon)
            coeffs->coeff[coeffs->number++].delay = i;
        else if (msg)
            fprintf(msg, " (less than epsilon, skipped)");
        if (msg) fprintf(msg, "\n");                    /* TAKE CARE FOR NO COEFFS AT ALL! */
        }
    if (msg) fprintf(msg, "\n");
    free(idx);                    /* Sure it's there. */
    return 0;

}

/*----------------------------------------------------------------------------*/
int COEFF_calc(COEFF** ab, POZE* pz, FILE* msg)
{
    int e;

    *ab = malloc(sizeof(COEFF));
    if (*ab)
        {
        if (msg)
            fprintf(msg, "%d zeros:\n", pz->zero.number);
        e = multiply_roots(&pz->zero, &(*ab)->a, msg);
        if (msg)
            fprintf(msg, "%d poles:\n", pz->pole.number);
        e |= multiply_roots(&pz->pole, &(*ab)->b, msg); /* Always do this init! */
        if (e)
            COEFF_free(ab);     /* Release possibly allocated members. */
        else                    /* Compensate A coefficients for gain, */
            {                   /* and also for time-displacement.     */
            int                 i;
            complexCartesian    rect_gain;

            polarToCartesian(&pz->gain, &rect_gain);
            for (i = 0; i < (*ab)->a.number; i++)
                {
                cartesianMultiply(&rect_gain, &(*ab)->a.coeff[i].cartesian);
                (*ab)->a.coeff[i].delay += (pz->pole.number - pz->zero.number);
                }                           /* Non-causal when this is > 0. */
            }
        }
    else
        e = 1;                  /* Not enough memory for stuct. */
    if (e && msg)
        fprintf(msg, "Memory failure in COEFF_calc()!\n");
    return e;
}

/*------------------------------------------------------------------------------*/
void COEFF_free(COEFF** ab)
{
    if (ab && *ab)
        {
        if ((*ab)->a.coeff)
            free((*ab)->a.coeff);
        if ((*ab)->b.coeff)
            free((*ab)->b.coeff);
        free(*ab);
        *ab = (COEFF*)NULL;
        }
}

/*----------------------------------------------------------------------------*/
int COEFF_write(const COEFF* ab, const char* stamp, FILE* fp)
{
    char    buffy[128];
    int     i;

    if (!(ab && fp))
        return 1;                       /* Bad arguments. */
    if (stamp)
        fprintf(fp, "# %s\n", stamp);
    for (i = 0; i < ab->a.number; i++)
        {
        writeCartesian(&ab->a.coeff[i].cartesian, 12, buffy);
        fprintf(fp, " a(%d) = %s\n", ab->a.coeff[i].delay, buffy);
        }
    fprintf(fp, "\n");
    for (i = 0; i < ab->b.number; i++)
        {
        writeCartesian(&ab->b.coeff[i].cartesian, 12, buffy);
        fprintf(fp, " b(%d) = %s\n", ab->b.coeff[i].delay, buffy);
        }
    fprintf(fp, "\n");
    return 0;
}

/*----------------------------------------------------------------------------*/
/*
    Same as COEFF_write() but now write as difference equation.
*/
int COEFF_write_diff_eq(const COEFF* ab, FILE* fp)
{
    char    buffy[128], *r;
    int     i, k, d;
    short   both;

    if (!(ab && fp))
        return 1;                   /* Bad arguments. */

    k = ab->a.number + ab->b.number;
    if      (k <  4) d = 8;
    if      (k <  5) d = 6;         /* Adapt number of decimals. */
    else if (k <  6) d = 5;
    else if (k <  8) d = 4;
    else if (k < 10) d = 3;
    else             d = 2;
    
    fprintf(fp, "Linear ");
    if (ab->a.coeff[0].delay < 0)   /* Hopefully this coeff exists! */
        fprintf(fp, "(non-causal) ");
    fprintf(fp, "difference equation:\n\n");

    for (i = 0; i < ab->b.number; i++)
        {                                      /* - means cut away trailing zeros. */
        writeCartesian(&ab->b.coeff[i].cartesian, -d, buffy);
        fprintf(fp, " ");
        if ((fabs(ab->b.coeff[i].cartesian.re) >= kPZT_epsilon) &&
            (fabs(ab->b.coeff[i].cartesian.im) >= kPZT_epsilon))
            {
            if (i)
                fprintf(fp, "+ ");
            fprintf(fp, "(");   /* Open bracket for complex coefficients. */
            both = 1;
            }
        else
            {
            if (i && (buffy[0] != '-'))
                fprintf(fp, "+ ");
            both = 0;
            }
        r = buffy;
        if (*r == '-')
            { fprintf(fp, "-"); r++; }
        if (strcmp(r, "1"))
            fprintf(fp, "%s", r);
        if (both)
            fprintf(fp, ")");   /* Close bracket for complex coefficients. */
        fprintf(fp, " y[n");
        if (ab->b.coeff[i].delay)
            fprintf(fp, "-%d", ab->b.coeff[i].delay);
        fprintf(fp, "]");
        }
    fprintf(fp, " =");
    for (i = 0; i < ab->a.number; i++)
        {
        writeCartesian(&ab->a.coeff[i].cartesian, -d, buffy);
        fprintf(fp, " ");
        if ((fabs(ab->a.coeff[i].cartesian.re) >= kPZT_epsilon) &&
            (fabs(ab->a.coeff[i].cartesian.im) >= kPZT_epsilon))
            {
            if (i)
                fprintf(fp, "+ ");
            fprintf(fp, "(");   /* Open bracket for complex coefficients. */
            both = 1;
            }
        else
            {
            if (i && (buffy[0] != '-'))
                fprintf(fp, "+ ");
            both = 0;
            }
        r = buffy;
        if (*r == '-')
            { fprintf(fp, "-"); r++; }
        if (strcmp(r, "1"))
            fprintf(fp, "%s", r);
        if (both)
            fprintf(fp, ")");   /* Close bracket for complex coefficients. */
        fprintf(fp, " x[n");
        if (ab->a.coeff[i].delay)
            {
            if (ab->a.coeff[i].delay > 0)
                fprintf(fp, "-%d", ab->a.coeff[i].delay);
            else
                fprintf(fp, "+%d", -ab->a.coeff[i].delay);
            }
        fprintf(fp, "]");
        }
    fprintf(fp, "\n");
    return 0;
}


/*----------------------------------------------------------------------------*/
int COEFF_read(COEFF**     ab,
               short       content, /* 0 if next arg=file. */
               const char* source,  /* filename or content */
               char*       msg)
{
    static const char*  errstr = " ERROR in COEFF_read(): ";
    static const int    maxLineLength = 256;   /* Including the \0 character. */
    FILE                *f = (FILE*)NULL;
    char                *lineBuff, *is;
    const char*         mem_src;
    cartesian_coeff     *a = NULL,      /* Write pointers. */
                        *b = NULL;
    complexCartesian    recRect;        /* To receive complex numbers in both */
    complexPolar        recPolar;       /* cartesian and polar format.        */
    int                 d, len,
                        e = 0,
                        linenum = 0;
    *ab = NULL;
    if (!msg)
        return 1;
    if (!source)
        { sprintf(msg, "%sthird argument NULL!\n", errstr); return 2; }
    if (!content)
        {
        f = fopen(source, "rt");            /* Argument preferably not NULL. */
        if (!f)
            { sprintf(msg, "%s'%s' not found!\n", errstr, source); return 3; }
        }
    else
        mem_src = source;
    lineBuff = malloc(maxLineLength);
    if (!lineBuff)
        { if (f) fclose(f); sprintf(msg, "%snot enough memory!\n", errstr); return 4; }
    *ab = malloc(sizeof(COEFF));
    if (!*ab)
        { free(lineBuff); if (f) fclose(f); sprintf(msg, "%snot enough memory!\n", errstr); return 5; }

    (*ab)->a.number = 0;                /* First count poles + zeros. */
    (*ab)->b.number = 0;
    (*ab)->a.coeff = NULL;
    (*ab)->b.coeff = NULL;
    while ((  content  && copy_line(lineBuff, maxLineLength, &mem_src)) ||
           ((!content) && fgets(lineBuff, maxLineLength, f)))
        {
        linenum++;
        len = simplifyString(lineBuff);            
        if (len && (lineBuff[0] != '#'))    /* Ignore empty lines and     */
            {                               /* lines that start with '#'. */
            if (!strchr("ab", lineBuff[0]))
                {
                sprintf(msg, "%sLine %d doesn't start with 'a' or 'b'!:\n %s\n", errstr, linenum, lineBuff);
                e = 6; goto done;
                }
            is = strchr(lineBuff + 1, '=');
            if (!is)
                {
                sprintf(msg, "%sNo '=' found in line %d!:\n %s\n", errstr, linenum, lineBuff);
                e = 7; goto done;
                }
            if ((lineBuff[1] != '(') || (is[-1] != ')'))
                {
                sprintf(msg, "%sMissing or misbalanced brackets in line %d!:\n %s\n", errstr, linenum, lineBuff);
                e = 8; goto done;
                }
            if (is == lineBuff + 2)
                {
                sprintf(msg, "%sEmpty brackets in line %d!:\n %s\n", errstr, linenum, lineBuff);
                e = 9; goto done;
                }
            if (lineBuff[0] == 'a') (*ab)->a.number++;
            else                    (*ab)->b.number++;
            }
        }
    if (!((*ab)->a.number && (*ab)->b.number))
        { sprintf(msg, "%sThere must at least be 1 a coeff and 1 b coeff!\n", errstr); e = 10; goto done; }
    a = (*ab)->a.coeff = (cartesian_coeff*)malloc((*ab)->a.number * sizeof(cartesian_coeff));
    b = (*ab)->b.coeff = (cartesian_coeff*)malloc((*ab)->b.number * sizeof(cartesian_coeff));
    if (!(a && b))
        { sprintf(msg, "%snot enough memory for a coeffs!\n", errstr); e = 11; goto done; }

    if (content)
        mem_src = source;
    else
        rewind(f);
    linenum = 0;                              /* Now really load them.  */
    while ((  content  && copy_line(lineBuff, maxLineLength, &mem_src)) ||
           ((!content) && fgets(lineBuff, maxLineLength, f)))            
        {
        linenum++;
        simplifyString(lineBuff);
        if (lineBuff[0] && (lineBuff[0] != '#'))
            {
            is = strchr(lineBuff + 1, '='); /* Sure we'll find it. */
            is[-1] = 0;                     /* Overwrite closing bracket. */
            d = atoi(lineBuff + 2);
            if (readPolarAndCartesian(is + 1, &recPolar, &recRect))
                {
                sprintf(msg, "%sInvalid polar or cartesian notation in line %d:\n %s\n",
                             errstr, linenum, lineBuff);
                e = 12; break;
                }
            if (*lineBuff == 'a')
                {
                a->delay = d; 
                memcpy(&a->cartesian, &recRect, sizeof(complexCartesian));
                a++;
                }
            else
                {
                if (d < 0)
                    {
                    sprintf(msg, "%sB indices may not be negative!\n", errstr);
                    e = 13; break;
                    }
                b->delay = d;
                memcpy(&b->cartesian, &recRect, sizeof(complexCartesian));
                b++;
                }
            }
        }
done:
    if (e)
        COEFF_free(ab);
    else
        sprintf(msg, " Read %d a-coefficients and %d b-coefficients.\n", (*ab)->a.number, (*ab)->b.number);
    free(lineBuff); /* Sure it's there. */
    if (f)
        fclose(f);
    return e;
}