#include <jerasure/reed_sol.h>
#include <jerasure.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

// Generate LRC matrix: (groups*local + global) code rows with (data_drives) columns
// w should be >= log2(data_drives + groups*local + global), but not necessary 8/16/32
int* reed_sol_vandermonde_lrc_matrix(int data_drives, int groups, int local, int global, int w)
{
    if (w < 0 || w > 32 || data_drives + groups*local + global > (1<<w))
    {
        return NULL;
    }
    int *lrc_matrix = (int*)malloc(sizeof(int) * (local*groups+global));
    int *matrix = reed_sol_vandermonde_coding_matrix(data_drives, local+global, w);
    // Enough to transform LRC 8+2+2 GF(8) matrix into MR-LRC
    //for (int i = 0; i < local+global; i++)
    //{
    //    int t = matrix[i*data_drives + 3];
    //    matrix[i*data_drives + 3] = matrix[i*data_drives + 7];
    //    matrix[i*data_drives + 7] = t;
    //}
    for (int gr = 0; gr < groups; gr++)
    {
        for (int l = 0; l < local; l++)
        {
            for (int j = 0; j < data_drives; j++)
            {
                lrc_matrix[(gr*local+l)*data_drives + j] = (j / (data_drives/groups)) == gr ? matrix[l*data_drives + j] : 0;
            }
        }
    }
    for (int i = 0; i < global; i++)
    {
        for (int j = 0; j < data_drives; j++)
        {
            lrc_matrix[(groups*local+i)*data_drives + j] = matrix[(local+i)*data_drives + j];
        }
    }
    free(matrix);
    return lrc_matrix;
}

struct lrc_test_result_t
{
    int success, impossible, failures;
};

// Check if the generated LRC with given parameters is Maximally Reconstructible (MR-LRC)
// Example of a MR-LRC: (8, 2, 1, 2, 6, 8)
struct lrc_test_result_t check_mr_lrc(int *lrc_matrix, int data_drives, int groups, int local, int global, int w, int log_level)
{
    int n = data_drives;
    int total_rows = n + groups*local + global;
    int impossible = 0, success = 0, failures = 0;
    int *lost_per_group = (int*)malloc(sizeof(int) * groups);
    int *recovered_per_group = (int*)malloc(sizeof(int) * groups);
    int *selected_inverted = (int*)malloc(sizeof(int) * data_drives);
    // global+1 is always recoverable
    for (int lost = global+2; lost <= groups*local+global; lost++)
    {
        int *erased_matrix = (int*)malloc(sizeof(int) * (total_rows-lost)*n);
        int *inverted_matrix = (int*)malloc(sizeof(int) * (total_rows-lost)*n);
        int *p = (int*)malloc(sizeof(int) * (total_rows-lost));
        for (int i = 0; i < n; i++)
            p[i] = i;
        int *p2 = (int*)malloc(sizeof(int) * n);
        if (total_rows-lost > n)
        {
            p[n-1] = n; // skip combinations with all N data disks (0..n-1)
            for (int i = n; i < total_rows-lost; i++)
                p[i] = i+1;
            p[total_rows-lost-1]--; // will be incremented on the first step
        }
        int inc = total_rows-lost-1;
        while (1)
        {
            p[inc]++;
            if (p[inc] >= n+groups*local+global)
            {
                if (inc == 0)
                    break;
                inc--;
            }
            else if (inc+1 < total_rows-lost)
            {
                p[inc+1] = p[inc];
                inc++;
            }
            else
            {
                // Check if it should be recoverable
                // Calculate count of data chunks lost in each group
                int nsel = 0;
                for (int gr = 0; gr < groups; gr++)
                {
                    lost_per_group[gr] = ((gr+1)*(n/groups) > n ? (n - gr*(n/groups)) : n/groups);
                    recovered_per_group[gr] = 0;
                }
                for (int j = 0; j < total_rows-lost; j++)
                {
                    if (p[j] < n)
                    {
                        lost_per_group[(p[j] / (n/groups))]--;
                        selected_inverted[nsel++] = j;
                    }
                }
                // Every local parity chunk is supposed to restore 1 missing chunk inside its group
                // So, subtract local parity chunk counts from each group lost chunk count
                for (int j = 0; j < total_rows-lost; j++)
                {
                    if (p[j] >= n && p[j] < n+groups*local)
                    {
                        int gr = (p[j]-n)/local;
                        if (lost_per_group[gr] > recovered_per_group[gr] && nsel < n)
                        {
                            selected_inverted[nsel++] = j;
                        }
                        recovered_per_group[gr]++;
                    }
                }
                // Every global parity chunk is supposed to restore 1 chunk of all that are still missing
                int still_missing = 0;
                for (int gr = 0; gr < groups; gr++)
                {
                    int non_fixed = lost_per_group[gr] - recovered_per_group[gr];
                    still_missing += (non_fixed > 0 ? non_fixed : 0);
                }
                for (int j = 0; j < total_rows-lost; j++)
                {
                    if (p[j] >= n+groups*local)
                    {
                        if (still_missing > 0 && nsel < n)
                        {
                            selected_inverted[nsel++] = j;
                        }
                        still_missing--;
                    }
                }
                if (still_missing <= 0)
                {
                    // We hope it can be recoverable. Try to invert it
                    assert(nsel == n);
                    for (int i = 0; i < n; i++)
                    {
                        for (int j = 0; j < n; j++)
                        {
                            erased_matrix[i*n+j] = lrc_matrix[p[selected_inverted[i]]*n+j];
                        }
                    }
                    int invert_ok = jerasure_invert_matrix(erased_matrix, inverted_matrix, n, w);
                    if (invert_ok < 0)
                    {
                        failures++;
                        if (log_level > 0)
                        {
                            printf("\nFAIL: ");
                            for (int i = 0; i < total_rows-lost; i++)
                            {
                                printf("%d ", p[i]);
                            }
                            printf("\nDIRECT:\n");
                            for (int i = 0; i < total_rows-lost; i++)
                            {
                                for (int j = 0; j < n; j++)
                                    printf("%d ", lrc_matrix[p[i]*n+j]);
                                printf("\n");
                            }
                            printf("INVERSE:\n");
                            for (int i = 0; i < total_rows-lost; i++)
                            {
                                for (int j = 0; j < n; j++)
                                    printf("%d ", inverted_matrix[i*n+j]);
                                printf("\n");
                            }
                        }
                    }
                    else
                    {
                        success++;
                        if (log_level > 2)
                        {
                            printf("OK: ");
                            for (int i = 0; i < total_rows-lost; i++)
                            {
                                printf("%d ", p[i]);
                            }
                            printf("\n");
                        }
                    }
                }
                else
                {
                    impossible++;
                    if (log_level > 1)
                    {
                        printf("IMPOSSIBLE: ");
                        for (int i = 0; i < total_rows-lost; i++)
                        {
                            printf("%d ", p[i]);
                        }
                        printf("\n");
                    }
                }
            }
        }
        free(p2);
        free(p);
        free(inverted_matrix);
        free(erased_matrix);
    }
    free(lost_per_group);
    free(recovered_per_group);
    return (struct lrc_test_result_t){
        .success = success,
        .impossible = impossible,
        .failures = failures,
    };
}

int main()
{
    int W = 8, MATRIX_W = 8;
    int n = 8, groups = 2, local = 1, global = 2;
    //n = 4, groups = 2, local = 1, global = 1;
    int total_rows = n+groups*local+global;
    int *matrix = reed_sol_vandermonde_lrc_matrix(n, groups, local, global, MATRIX_W);
    int *lrc_matrix = (int*)malloc(sizeof(int) * total_rows*n);
    // Fill identity+LRC matrix
    for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++)
            lrc_matrix[i*n + j] = j == i ? 1 : 0;
    memcpy(lrc_matrix + n*n, matrix, (total_rows-n)*n*sizeof(int));
    free(matrix);
    matrix = NULL;
    // Print LRC matrix
    for (int i = 0; i < total_rows; i++)
    {
        for (int j = 0; j < n; j++)
        {
            printf("%d ", lrc_matrix[i*n+j]);
        }
        printf("\n");
    }
    struct lrc_test_result_t t = check_mr_lrc(lrc_matrix, n, groups, local, global, W, 1);
    printf("\n%d recovered, %d impossible, %d failures\n", t.success, t.impossible, t.failures);
    return 0;
}

// 1 1 1 1 0 0 0 0
// 0 0 0 0 1 1 1 1
// 1 55 39 73 84 181 225 217
// 1 172 70 235 143 34 200 101
//
// Can't recover
// 1 2 4 5 8 9 10 11    -1
// 2 3 4 6 8 9 10 11    -1
// FULL:
// 1 0 0 0 0 0 0 0 
// 0 1 0 0 0 0 0 0 
// 0 0 1 0 0 0 0 0 
// 0 0 0 1 0 0 0 0 
// 0 0 0 0 1 0 0 0 
// 0 0 0 0 0 1 0 0 
// 0 0 0 0 0 0 1 0 
// 0 0 0 0 0 0 0 1 
// 1 1 1 1 0 0 0 0 
// 0 0 0 0 1 1 1 1 
// 1 55 39 73 84 181 225 217 
// 1 172 70 235 143 34 200 101 
// FIRST UNRECOVERABLE:
// 0 1 0 0 0 0 0 0 
// 0 0 1 0 0 0 0 0 
// 0 0 0 0 1 0 0 0 
// 0 0 0 0 0 1 0 0 
// 1 1 1 1 0 0 0 0 
// 0 0 0 0 1 1 1 1 
// 1 55 39 73 84 181 225 217 
// 1 172 70 235 143 34 200 101 
// SECOND UNRECOVERABLE:
// 0 0 1 0 0 0 0 0 
// 0 0 0 1 0 0 0 0 
// 0 0 0 0 1 0 0 0 
// 0 0 0 0 0 0 1 0 
// 1 1 1 1 0 0 0 0 
// 0 0 0 0 1 1 1 1 
// 1 55 39 73 84 181 225 217 
// 1 172 70 235 143 34 200 101 
// Ho ho ho