#include "ncnf.h"

cnf* readDIMACS(char* path) {
    cnf* c = malloc(sizeof(cnf));
    CHECK(c, "Failed to alloc CNF struct\n")

    FILE* f = fopen(path, "r");
    CHECK(f, "Failed to open file\n")

    u32 bufsize = 1000;
    char* buf = malloc(sizeof(char) * bufsize);
    CHECK(buf, "Failed to alloc read buffer\n")

    CHECK(fgets(buf, sizeof(char) * bufsize, f), "Failed to read file\n")

    // 'c' marks the beginning of a comment line in DIMACs. Skip this line
    while (buf[0] == 'c') {
        CHECK(fgets(buf, sizeof(char) * bufsize, f), "Failed to read file\n")
    }

    // The first non-comment line is the header, should have format 'p cnf [varcnt] [clausecnt]
    // We skip the first bit by iterating and skipping any character that's p, c, n, f, or space
    char* temp = buf;
    while (*temp == ' ' || *temp == 'p' || *temp == 'c' || *temp == 'n' || *temp == 'f') {
        temp++;
    }

    // Char by char we read in the number of variables
    u32* varcnt = &(c->cnts[0]);
    *varcnt = 0;
    while (((u8) (*temp - '0')) < 10) {
        *varcnt *= 10;
        *varcnt += (*temp - '0');
        temp++;
    }

    // Skip any trailing whitespace
    while (*temp == ' ') temp++;

    // Read in clausecnt
    u32* clausecnt = &(c->cnts[1]);
    *clausecnt = 0;
    while (((u8) (*temp - '0')) < 10) {
        *clausecnt *= 10;
        *clausecnt += (*temp - '0');
        temp++;
    }

    // Resize line buffer: the maximum clause size is varcnt, the maximum size of a u32 in decimal
    // is 10, and conservatively given 2 extra characters for the - symbol and space delimiting,
    // resizing the buffer to this should mean we never encounter a clause we fail to read.
    u32 nsize = sizeof(char) * ((*varcnt) * 12LU) + 10;
    char* nbuf = realloc(buf, nsize);
    CHECK(nbuf, "Failed to realloc read buffer\n")

    buf = nbuf;

    u32 cnt = 0;
    u32 cap = 32;

    c->clausedat = calloc(*clausecnt, sizeof(u32) * 3);
    CHECK(c->clausedat, "Failed to allocate clause data\n")
    c->index = calloc(*varcnt, sizeof(u32));
    CHECK(c->index, "Failed to allocate clause index\n")
    c->variables = calloc(cap, sizeof(u32));
    CHECK(c->variables, "Failed to allocate literal variables\n")
    c->parities = calloc(cap, sizeof(u8));
    CHECK(c->parities, "Failed to allocate literal parities\n");

    for (u32 i = 0; i < *clausecnt; ++i) {
        c->clausedat[3 * i] = cnt;
        CHECK(fgets(buf, nsize, f), "Failed to read clause\n");

        temp = buf;
        while (*temp == ' ') temp++;

        bool empty = true;
        bool tr = true;
        while (*temp != '\n') {
            if (cnt == cap) {
                // Out of space to add more literals, realloc arrays
                u32 ncap = cap << 1U;
                c->variables = realloc(c->variables, sizeof(u32) * ncap);
                memset(c->variables + cap, 0, sizeof(u32) * cap);
                CHECK(c->variables, "Failed to realloc variable array\n");
                c->parities = realloc(c->parities, sizeof(u8) * ncap);
                memset(c->parities + cap, 0, sizeof(u8) * cap);
                CHECK(c->parities, "Failed to realloc parity array\n");
                cap = ncap;
            }

            if (*temp == '-') {
                tr = false;
            } else if (((u8) (*temp - '0'))  < 10) {
                // Read in the literal's digits
                c->variables[cnt] *= 10;
                c->variables[cnt] += (*temp - '0');
            } else {
                // Skip any whitespace and pack the read value into the arrays, along with
                // any additional data
                while (temp[1] == ' ') temp++;
                c->parities[cnt] = tr;
                c->variables[cnt] -= 1;
                if (c->clausedat[3 * i + 2] < c->variables[cnt]) c->clausedat[3 * i + 2] = c->variables[cnt];
                empty = false;
                tr = true;
                cnt++;
                if (temp[1] == '0') {
                    break;
                }
            }
            temp++;
        }

        if (empty) {
            printf("UNSAT: Empty clause %u\n", i);
            return NULL;
        }
        c->clausedat[3 * i + 1] = cnt - c->clausedat[3 * i];
    }

    // Lastvars is set to the index of the last variable of each clause. However, because in the
    // counter, the last variable has the lsb position, we flip the value of lastvars
    for (u32 i = 0; i < *clausecnt; ++i) {
        c->clausedat[3 * i + 2] += 1U;
        c->clausedat[3 * i + 2] = *varcnt - c->clausedat[3 * i + 2];
    }

    // Realloc the arrays to exactly match the number of literals
    c->variables = realloc(c->variables, sizeof(u32) * cnt);
    c->parities = realloc(c->parities, sizeof(u8) * cnt);
    c->cnts[2] = cnt;
    free(buf);
    if (fclose(f)) {
        printf("Failed to close file\n");
        return NULL;
    }
    return c;
}

void printcnf(cnf* c) {
    printf("p cnf %u %u\n", c->cnts[0], c->cnts[1]);
    for (u32 i = 0; i < c->cnts[1]; ++i) {
        for (u32 j = 0; j < c->clausedat[3 * i + 1]; ++j) {
            if (c->parities[c->clausedat[3 * i] + j] == 0) printf("-");
            printf("%u ", c->variables[c->clausedat[3 * i] + j] + 1);
        }
        printf("0\n");
    }
}

void freecnf(cnf* c) {
    free(c->clausedat);
    free(c->variables);
    free(c->parities);
    free(c);
}

const u64 RBITS = 8;
const u64 RSIZE = 1LU << RBITS;
const u64 RLVLS = (31LU / RBITS) + 1;
const u64 RMASK = RSIZE - 1;

void sortlastnum(cnf* c) {
    bool v = false;
    u32* d = calloc(c->cnts[1], sizeof(u32) * 3);

    u32 qb[RLVLS * RSIZE];
    u32 qe[RLVLS * RSIZE];
    memset(qb, 0, sizeof(u32) * RLVLS * RSIZE);
    memset(qe, 0, sizeof(u32) * RLVLS * RSIZE);

    for (u32 pass = 0; pass < RLVLS; ++pass) {
        u32 shift = pass * RBITS;

        for (u32 i = 0; i < c->cnts[1]; ++i) {
            u32 val = UINT32_MAX - c->clausedat[3 * i + 2];
            u32 ind = (val >> shift) & RMASK;
            qe[pass * RSIZE + ind]++;
        }

        u32 uc = 0;
        for (u32 i = 0; i < RSIZE; ++i) if (qe[pass * RSIZE + i]) uc++;
        if (uc == 1) continue;

        qb[pass * RSIZE] = 0;
        for (u32 i = 1; i < RSIZE; ++i) {
            qb[pass * RSIZE + i] = qb[pass * RSIZE + i - 1] + qe[pass * RSIZE + i - 1];
        }

        for (u32 i = 0; i < c->cnts[1]; ++i) {
            u32 val = UINT32_MAX - c->clausedat[3 * i + 2];
            u32 ind = (val >> shift) & RMASK;
            memcpy(d + 3 * qb[pass * RSIZE + ind], c->clausedat + 3 * i, sizeof(u32) * 3);
            qb[pass * RSIZE + ind]++;
            __builtin_prefetch(d + 3 * qb[pass * RSIZE + ind] + 1);
        }

        u32* tptr = c->clausedat;
        c->clausedat = d;
        d = tptr;
        v = !v;
    }

    if (v) {
        u32* tptr = c->clausedat;
        c->clausedat = d;
        d = tptr;
        memcpy(c->clausedat, d, sizeof(u32) * 3 * c->cnts[1]);
    }

    free(d);

    // TODO: Rewrite, the following is copied from Gábor Kusper's implementation
    u32 lastNum = 0;
    for (u32 i = c->cnts[1] - 1; i < c->cnts[1]; --i) {
        if (c->clausedat[3 * i + 2] > lastNum) {
            while (lastNum < c->clausedat[3 * i + 2]) {
                c->index[lastNum] = i + 1;
                lastNum++;
            }
        }
    }


    u32 corrInd = 0;
    while (c->index[corrInd] == c->cnts[1]) {
        corrInd++;
    }

    if (corrInd != 0) {
        u32 goodVal = c->index[corrInd];
        do {
            corrInd--;
            c->index[corrInd] = goodVal;
        } while (corrInd != 0);
    }
}