#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <sys/time.h>
#include "../ytools/ytools.h"
#include "batch_factor.h"
#include "gmp.h"

// build line:
// clang -O2 -g -I. -I../../ytools/ -I../../include -I../../ms_include -I../../top/aprcl -I../../../gmp-install/6.2.0-aocc/include -L../../../gmp-install/6.2.0-aocc/lib -march=icelake-client -DUSE_AVX512F -DIFMA ../../ytools/util.c batch_factor.c process_batch.c tinyecm.c microecm.c prime_sieve.c micropm1.c -o bfact -lm -lgmp
// 
// usage: bfact lpb pmin pmax relsfilein relsfileout

int main(int argc, char** argv)
{
    relation_batch_t rb;

    uint64_t lpb = atoi(argv[1]);
    uint64_t pmin = strtoull(argv[2], NULL, 10);
    uint64_t pmax = strtoull(argv[3], NULL, 10);
    char infile[80], outfile[80];
    char buf[1024], str1[1024], str2[1024];
    uint32_t fr[32], fa[32], numr = 0, numa = 0;
    mpz_t res1, res2;
    struct timeval start;
    struct timeval stop;
    double ttime;
    uint64_t lcg_state = 42;
    int i;
    uint32_t line = 0;
    uint32_t numfull = 0;

    strcpy(infile, argv[4]);
    strcpy(outfile, argv[5]);

    if (lpb > 40)
    {
        printf("expected large prime bound < 40 bits, saw %lu\n", lpb);
        exit(0);
    }

    lpb = 1ull << lpb;

    for (i = 0; i < 4; i++)
    {
        rb.num_uecm[i] = 0;
        rb.num_uecm_a[i] = 0;
    }
    rb.num_tecm = 0;
    rb.num_tecm2 = 0;
    rb.num_qs = 0;
    rb.num_tecm_a = 0;
    rb.num_tecm2_a = 0;
    rb.num_qs_a = 0;
    rb.num_attempt = 0;
    rb.num_success = 0;
    for (i = 0; i < 8; i++)
    {
        rb.num_abort[i] = 0;
        rb.num_abort_a[i] = 0;
    }
    mpz_init(res1);
    mpz_init(res2);

    printf("initializing relation batch...\n");

    gettimeofday(&start, NULL);
    relation_batch_init(stdout, &rb, pmin, pmax, lpb, lpb, NULL, 1);
    gettimeofday(&stop, NULL);
    ttime = ytools_difftime(&start, &stop);

    printf("init took %1.2f sec, reading input file...\n", ttime);

    FILE* fid = fopen(infile, "r");
    if (fid == NULL)
    {
        printf("could not open %s to read\n", infile);
        exit(0);
    }

    FILE* fout; // = fopen(outfile, "w");
    //if (fid == NULL)
    //{
    //    printf("could not open %s to write\n", outfile);
    //    exit(0);
    //}

    gettimeofday(&start, NULL);

    while (~feof(fid))
    {
        int64_t a;
        uint32_t b;
        char* tok;

        line++;
        char *ptr = fgets(buf, 1024, fid);
        if (ptr == NULL)
            break;

        // printf("buffer read: %s\n", buf);
        strcpy(str1, buf);

        tok = strtok(buf, ":");
        if (tok == NULL)
        {
            printf("could not read relation %u, no lfactors token\n", line);
            continue;
        }

        ptr = strchr(tok, ',');
        *ptr = '\0';
        mpz_set_str(res1, tok, 10);
        mpz_set_str(res2, ptr+1, 10);

        // gmp_printf("large factors: %Zd, %Zd\n", res1, res2);

        tok = strtok(NULL, ":");
        if (tok == NULL)
        {
            printf("could not read relation %u, no a/b token\n", line);
            continue;
        }

        sscanf(tok, "%ld,%u", &a, &b);

        tok = strtok(NULL, ":");
        if (tok == NULL)
        {
            printf("could not read relation %u, no rfactors token\n", line);
            continue;
        }
        numr = 0;
        while (strlen(tok) > 0)
        {
            fr[numr++] = strtoul(tok, NULL, 16);
            ptr = strchr(tok, ',');
            if (ptr == NULL)
                break;
            tok = ptr + 1;
        }

        tok = strtok(NULL, ":");
        if (tok == NULL)
        {
            printf("could not read relation %u, no afactors token\n", line);
            continue;
        }

        numa = 0;
        while (strlen(tok) > 0)
        {
            fa[numa++] = strtoul(tok, NULL, 16);
            ptr = strchr(tok, ',');
            if (ptr == NULL)
                break;
            tok = ptr + 1;
        }

        if ((mpz_sgn(res1) > 0) && (mpz_sgn(res2) > 0))
        {
            numfull++;
            //fprintf(fout, "%s", str1);
        }
        else
        {
            relation_batch_add(a, b, fr, numr, res1, fa, numa, res2, &rb);
        }
    }
    //fclose(fout);
    fclose(fid);
    //printf("wrote %d full relations to %s\n", numfull, outfile);

    gettimeofday(&stop, NULL);
    ttime = ytools_difftime(&start, &stop);
    printf("file parsing took %1.2f sec, found %d fulls, batched %u rels, now running batch solve...\n", 
        ttime, numfull, rb.num_relations);

    gettimeofday(&start, NULL);
    relation_batch_run(&rb, &lcg_state);
    gettimeofday(&stop, NULL);

    ttime = ytools_difftime(&start, &stop);
    printf("relation_batch_run took %1.4f sec producing %u relations\n",
        ttime, rb.num_success);

    fid = fopen(infile, "r");
    if (fid == NULL)
    {
        printf("could not open %s to read\n", infile);
        exit(0);
    }

    fout = fopen(outfile, "w");

    if (fout != NULL)
    {
        int num_success = 0;
        int n = 0;
        int nwrote = 0;
        line = 0;
        for (i = 0; i < rb.num_relations; i++)
        {
            if (rb.relations[i].success > 0)
            {
                int j, k;

                uint32_t* f = rb.factors + rb.relations[i].factor_list_word;

                fprintf(fout, "%ld,%u:", rb.relations[i].a, rb.relations[i].b);
                for (j = 0; j < 3; j++)
                {
                    if (rb.relations[i].lp_r[j] > 1)
                        fprintf(fout, "%x,", rb.relations[i].lp_r[j]);
                }
                for (k = 0; k < rb.relations[i].num_factors_r - 1; k++)
                {
                    fprintf(fout, "%x,", f[k]);
                }
                fprintf(fout, "%x:", f[k]);
                for (j = 0; j < 3; j++)
                {
                    if (rb.relations[i].lp_a[j] > 1) 
                        fprintf(fout, "%x,", rb.relations[i].lp_a[j]);
                }

                f = rb.factors + rb.relations[i].factor_list_word + rb.relations[i].num_factors_r;
                for (k = 0; k < rb.relations[i].num_factors_a - 1; k++)
                {
                    fprintf(fout, "%x,", f[k]);
                }
                fprintf(fout, "%x\n", f[k]);
                nwrote++;
            }
        }
        fclose(fout);
        printf("wrote %d relations to %s\n", nwrote, outfile);
        fclose(fid);
    }
    else
    {
        printf("could not open %s to append\n", outfile);
    }

    if (0)
    {
        if (fout != NULL)
        {
            int num_success = 0;
            int n = 0;
            int nwrote = 0;
            line = 0;
            for (i = 0; i < rb.num_relations; i++)
            {
                if (rb.relations[i].success > 0)
                {
                    num_success++;
                    int p;

                    // todo: find line 'i' of the input file and output it with the large cofactors
                    // replaced with their factorizations.
                    qsort(rb.relations[i].lp_r, 3, sizeof(uint32_t), &qcomp_uint32);
                    qsort(rb.relations[i].lp_a, 3, sizeof(uint32_t), &qcomp_uint32);

                    //printf("searching file for relation %u, currently on line %u\n", i, n);
                    while (~feof(fid))
                    {
                        int64_t a;
                        uint32_t b;
                        char* tok;
                        char rfactors[1024];
                        char afactors[1024];

                        line++;
                        char* ptr = fgets(buf, 1024, fid);
                        if (ptr == NULL)
                            break;

                        // printf("buffer read: %s\n", buf);
                        buf[strlen(buf) - 1] = '\0';
                        strcpy(str1, buf);

                        tok = strtok(buf, ":");
                        if (tok == NULL)
                        {
                            printf("could not read relation %u, no lfactors token\n", line);
                            continue;
                        }

                        ptr = strchr(tok, ',');
                        *ptr = '\0';
                        mpz_set_str(res1, tok, 10);
                        mpz_set_str(res2, ptr + 1, 10);

                        // gmp_printf("large factors: %Zd, %Zd\n", res1, res2);

                        tok = strtok(NULL, ":");
                        if (tok == NULL)
                        {
                            printf("could not read relation %u, no a/b token\n", line);
                            continue;
                        }

                        sscanf(tok, "%ld,%u", &a, &b);

                        tok = strtok(NULL, ":");
                        if (tok == NULL)
                        {
                            printf("could not read relation %u, no rfactors token\n", line);
                            continue;
                        }

                        strcpy(rfactors, tok);

                        numr = 0;
                        while (strlen(tok) > 0)
                        {
                            fr[numr++] = strtoul(tok, NULL, 16);
                            ptr = strchr(tok, ',');
                            if (ptr == NULL)
                                break;
                            tok = ptr + 1;
                        }

                        tok = strtok(NULL, ":");
                        if (tok == NULL)
                        {
                            printf("could not read relation %u, no afactors token\n", line);
                            continue;
                        }

                        strcpy(afactors, tok);

                        numa = 0;
                        while (strlen(tok) > 0)
                        {
                            fa[numa++] = strtoul(tok, NULL, 16);
                            ptr = strchr(tok, ',');
                            if (ptr == NULL)
                                break;
                            tok = ptr + 1;
                        }

                        // A relation has been read.  If it is a full relation, 
                        // then output it, but then don't include it in the search 
                        // for this relation's line: these are added not to the batch solve.
                        if ((mpz_sgn(res1) > 0) && (mpz_sgn(res2) > 0))
                        {
                            fprintf(fout, "%ld,%u:", a, b);
                            if (mpz_cmp_ui(res1, 1) > 0)
                                gmp_fprintf(fout, "%Zx,", res1);
                            fprintf(fout, "%s:", rfactors);
                            if (mpz_cmp_ui(res2, 1) > 0)
                                gmp_fprintf(fout, "%Zx,", res2);
                            fprintf(fout, "%s\n", afactors);
                            nwrote++;
                            continue;
                        }


                        if (n == i)
                        {
                            // make sure this factorization applies to this relation
                            int good = 1;
                            int j;
                            int print_rres = 0;
                            int print_ares = 0;

                            if (mpz_sgn(res1) < 0)
                            {
                                mpz_neg(res1, res1);
                                print_rres = 0;
                            }
                            else if (mpz_cmp_ui(res1, 1) > 0)
                            {
                                print_rres = 1;
                            }

                            if (mpz_sgn(res2) < 0)
                            {
                                mpz_neg(res2, res2);
                                print_ares = 0;
                            }
                            else if (mpz_cmp_ui(res2, 1) > 0)
                            {
                                print_ares = 1;
                            }

                            for (j = 0; j < 3; j++)
                                good &= (mpz_tdiv_ui(res1, rb.relations[i].lp_r[j]) == 0);

                            for (j = 0; j < 3; j++)
                                good &= (mpz_tdiv_ui(res2, rb.relations[i].lp_a[j]) == 0);

                            if (!good)
                            {
                                printf("factors of relations array %d don't divide residues of "
                                    "input file line %d as expected\n", i, n);
                                gmp_printf("residues: %Zd, %Zd\n", res1, res2);
                                printf("rfactors: ");
                                for (j = 0; j < 3; j++)
                                    if (rb.relations[i].lp_r[j] > 1)
                                        printf("%u ", rb.relations[i].lp_r[j]);
                                printf("\nafactors: ");
                                for (j = 0; j < 3; j++)
                                    if (rb.relations[i].lp_a[j] > 1)
                                        printf("%lu ", rb.relations[i].lp_a[j]);
                                printf("\n");
                            }
                            else
                            {
                                fprintf(fout, "%ld,%u:", a, b);
                                if (print_rres)
                                {
                                    gmp_fprintf(fout, "%Zx,", res1);
                                }
                                else
                                {
                                    rb.relations[i].num_factors_r = 0;
                                    for (j = 0; j < 3; j++)
                                    {
                                        if (rb.relations[i].lp_r[j] > 1)
                                            fprintf(fout, "%x,", rb.relations[i].lp_r[j]);
                                    }
                                }
                                fprintf(fout, "%s:", rfactors);

                                if (print_ares)
                                {
                                    gmp_fprintf(fout, "%Zx,", res2);
                                }
                                else
                                {
                                    rb.relations[i].num_factors_a = 0;
                                    for (j = 0; j < 3; j++)
                                    {
                                        if (rb.relations[i].lp_a[j] > 1)
                                            fprintf(fout, "%x,", rb.relations[i].lp_a[j]);
                                    }
                                }
                                fprintf(fout, "%s\n", afactors);
                                nwrote++;
                            }
                            n++;
                            break;

                        }
                        n++;

                    }
                }
            }

            fclose(fout);
            printf("wrote %d relations to %s\n", nwrote, outfile);
            fclose(fid);
        }
        else
        {
            printf("could not open %s to append\n", outfile);
        }
    }

    printf("ECM stats R:\n");
    for (i = 0; i < 4; i++)
    {
        printf("%u;  ", rb.num_uecm[i]);
    }
    printf("%u;  ", rb.num_tecm);
    printf("%u;  ", rb.num_tecm2);
    printf("%u;  ", rb.num_qs);
    printf("\nECM stats A:\n");
    for (i = 0; i < 4; i++)
    {
        printf("%u;  ", rb.num_uecm_a[i]);
    }
    printf("%u;  ", rb.num_tecm_a);
    printf("%u;  ", rb.num_tecm2_a);
    printf("%u;  ", rb.num_qs_a);

    printf("\nAbort stats R:\n");
    for (i = 0; i < 8; i++)
    {
        printf("%u;  ", rb.num_abort[i]);
    }
    printf("\nAbort stats A:\n");
    for (i = 0; i < 8; i++)
    {
        printf("%u;  ", rb.num_abort_a[i]);
    }
    printf("\n");

    mpz_clear(res1);
    mpz_clear(res2);
    relation_batch_free(&rb);

    return 0;
}