#include <stdio.h>

#include "ruby.h"
#include "spglib.h"

VALUE Getspg = Qnil;
void Init_getspg(void);
VALUE method_getspg(VALUE self, VALUE r_size, VALUE r_lattice, VALUE r_position,
                    VALUE r_types, VALUE r_symprec, VALUE r_angle_symprec);
VALUE method_getptg(VALUE self, VALUE r_rotations);
VALUE method_find_primitive(VALUE self, VALUE r_lattice, VALUE r_position,
                            VALUE r_types, VALUE r_symprec,
                            VALUE r_angle_symprec);
VALUE method_get_dataset(VALUE self, VALUE r_lattice, VALUE r_position,
                         VALUE r_types, VALUE r_hall_num, VALUE r_symprec,
                         VALUE r_angle_symprec);
VALUE method_get_symmetry(VALUE self, VALUE r_lattice, VALUE r_position,
                          VALUE r_types, VALUE r_symprec,
                          VALUE r_angle_symprec);
VALUE method_standardize_cell(VALUE self, VALUE r_lattice, VALUE r_position,
                              VALUE r_types, VALUE r_to_primitive,
                              VALUE r_leave_distrtion, VALUE r_symprec,
                              VALUE r_angle_symprec);

void Init_getspg(void) {
    Getspg = rb_define_module("Getspg");
    rb_define_method(Getspg, "getptg", method_getptg, 1);
    rb_define_method(Getspg, "find_primitive", method_find_primitive, 5);
    rb_define_method(Getspg, "get_dataset", method_get_dataset, 6);
    rb_define_method(Getspg, "get_symmetry", method_get_symmetry, 5);
}

VALUE method_getptg(VALUE self, VALUE r_rotations) {
    int i, j, k, size, ptg_num;
    char symbol[6];
    int trans_mat[3][3];
    VALUE array, matrix, vector;

    size = RARRAY_LEN(r_rotations);
    int rotations[size][3][3];

    for (i = 0; i < size; i++) {
        for (j = 0; j < 3; j++) {
            for (k = 0; k < 3; k++) {
                rotations[i][j][k] = NUM2INT(rb_ary_entry(
                    rb_ary_entry(rb_ary_entry(r_rotations, i), j), k));
            }
        }
    }

    ptg_num = spg_get_pointgroup(symbol, trans_mat, rotations, size);
    array = rb_ary_new();
    rb_ary_push(array, rb_str_new2(symbol));
    rb_ary_push(array, INT2NUM(ptg_num));
    matrix = rb_ary_new();
    for (i = 0; i < 3; i++) {
        vector = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(vector, INT2NUM(trans_mat[i][j]));
        }
        rb_ary_push(matrix, vector);
    }
    rb_ary_push(array, matrix);

    return array;
}

VALUE method_find_primitive(VALUE self, VALUE r_lattice, VALUE r_position,
                            VALUE r_types, VALUE r_symprec,
                            VALUE r_angle_symprec) {
    int i, j, num_atom, num_prim_atom;
    double symprec, lattice[3][3];
    VALUE array, vector, lat_ary, pos_ary, typ_ary;

    num_atom = RARRAY_LEN(r_types);

    double position[num_atom][3];
    int types[num_atom];

    symprec = NUM2DBL(r_symprec);

    for (i = 0; i < num_atom; i++) {
        for (j = 0; j < 3; j++) {
            position[i][j] =
                NUM2DBL(rb_ary_entry(rb_ary_entry(r_position, i), j));
            types[i] = NUM2DBL(rb_ary_entry(r_types, i));
        }
    }

    for (i = 0; i < 3; i++) {
        for (j = 0; j < 3; j++) {
            lattice[i][j] =
                NUM2DBL(rb_ary_entry(rb_ary_entry(r_lattice, i), j));
        }
    }

    num_prim_atom =
        spgat_standardize_cell(lattice, position, types, num_atom, 1, 0,
                               symprec, NUM2DBL(r_angle_symprec));

    array = rb_ary_new();
    lat_ary = rb_ary_new();
    for (i = 0; i < 3; i++) {
        vector = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(vector, rb_float_new(lattice[i][j]));
        }
        rb_ary_push(lat_ary, vector);
    }
    rb_ary_push(array, lat_ary);

    pos_ary = rb_ary_new();
    typ_ary = rb_ary_new();
    for (i = 0; i < num_prim_atom; i++) {
        vector = rb_ary_new();
        rb_ary_push(typ_ary, INT2NUM(types[i]));
        for (j = 0; j < 3; j++) {
            rb_ary_push(vector, rb_float_new(position[i][j]));
        }
        rb_ary_push(pos_ary, vector);
    }

    rb_ary_push(array, pos_ary);
    rb_ary_push(array, typ_ary);

    return array;
}

VALUE method_get_dataset(VALUE self, VALUE r_lattice, VALUE r_position,
                         VALUE r_types, VALUE r_hall_num, VALUE r_symprec,
                         VALUE r_angle_symprec) {
    int i, j, k, num_atom;
    double symprec, lattice[3][3];
    SpglibDataset *dataset;
    VALUE mat, vec, row, array, r_tmat, r_oshift, r_rot, r_trans, r_wyckoffs,
        r_std_lattice, r_std_positions, r_std_types;

    num_atom = RARRAY_LEN(r_types);

    double position[num_atom][3];
    int types[num_atom];

    symprec = NUM2DBL(r_symprec);

    for (i = 0; i < num_atom; i++) {
        for (j = 0; j < 3; j++) {
            position[i][j] =
                NUM2DBL(rb_ary_entry(rb_ary_entry(r_position, i), j));
            types[i] = NUM2DBL(rb_ary_entry(r_types, i));
        }
    }

    for (i = 0; i < 3; i++) {
        for (j = 0; j < 3; j++) {
            lattice[i][j] =
                NUM2DBL(rb_ary_entry(rb_ary_entry(r_lattice, i), j));
        }
    }

    dataset = spgat_get_dataset_with_hall_number(
        lattice, position, types, num_atom, NUM2INT(r_hall_num), symprec,
        NUM2DBL(r_angle_symprec));

    array = rb_ary_new();

    if (dataset == NULL) {
        goto err;
    }

    rb_ary_push(array, INT2NUM(dataset->spacegroup_number));
    rb_ary_push(array, rb_str_new2(dataset->international_symbol));
    rb_ary_push(array, INT2NUM(dataset->hall_number));
    rb_ary_push(array, rb_str_new2(dataset->hall_symbol));
    rb_ary_push(array, rb_str_new2(dataset->choice));

    /* Transformation_matrix */
    r_tmat = rb_ary_new();
    for (i = 0; i < 3; i++) {
        row = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(row,
                        rb_float_new(dataset->transformation_matrix[i][j]));
        }
        rb_ary_push(r_tmat, row);
    }
    rb_ary_push(array, r_tmat);

    /* Origin shift */
    r_oshift = rb_ary_new();
    for (i = 0; i < 3; i++) {
        rb_ary_push(r_oshift, rb_float_new(dataset->origin_shift[i]));
    }
    rb_ary_push(array, r_oshift);

    /* Rotations, translations */
    r_rot = rb_ary_new();
    r_trans = rb_ary_new();
    for (i = 0; i < dataset->n_operations; i++) {
        mat = rb_ary_new();
        vec = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(vec, rb_float_new(dataset->translations[i][j]));
            row = rb_ary_new();
            for (k = 0; k < 3; k++) {
                rb_ary_push(row, rb_float_new(dataset->rotations[i][j][k]));
            }
            rb_ary_push(mat, row);
        }
        rb_ary_push(r_trans, vec);
        rb_ary_push(r_rot, mat);
    }
    rb_ary_push(array, r_rot);
    rb_ary_push(array, r_trans);

    /* Wyckoff letters */
    r_wyckoffs = rb_ary_new();
    for (i = 0; i < dataset->n_atoms; i++) {
        r_wyckoffs = rb_ary_push(r_wyckoffs, INT2NUM(dataset->wyckoffs[i]));
    }
    rb_ary_push(array, r_wyckoffs);

    /* Bravais lattice */
    r_std_lattice = rb_ary_new();
    for (i = 0; i < 3; i++) {
        vec = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(vec, rb_float_new(dataset->std_lattice[i][j]));
        }
        rb_ary_push(r_std_lattice, vec);
    }
    rb_ary_push(array, r_std_lattice);

    r_std_positions = rb_ary_new();
    r_std_types = rb_ary_new();
    for (i = 0; i < dataset->n_std_atoms; i++) {
        vec = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(vec, rb_float_new(dataset->std_positions[i][j]));
        }
        rb_ary_push(r_std_positions, vec);
        rb_ary_push(r_std_types, INT2NUM(dataset->std_types[i]));
    }
    rb_ary_push(array, r_std_types);
    rb_ary_push(array, r_std_positions);

    spg_free_dataset(dataset);

err:

    return array;
}

VALUE method_get_symmetry(VALUE self, VALUE r_lattice, VALUE r_position,
                          VALUE r_types, VALUE r_symprec,
                          VALUE r_angle_symprec) {
    int i, j, k, num_atom, num_sym;
    double symprec, lattice[3][3];
    VALUE mat, vec, row, array, r_rot, r_trans;

    num_atom = RARRAY_LEN(r_types);

    double position[num_atom][3];
    int types[num_atom];

    int max_num_sym = num_atom * 48;
    int rotations[max_num_sym][3][3];
    double translations[max_num_sym][3];

    symprec = NUM2DBL(r_symprec);

    for (i = 0; i < num_atom; i++) {
        for (j = 0; j < 3; j++) {
            position[i][j] =
                NUM2DBL(rb_ary_entry(rb_ary_entry(r_position, i), j));
            types[i] = NUM2DBL(rb_ary_entry(r_types, i));
        }
    }

    for (i = 0; i < 3; i++) {
        for (j = 0; j < 3; j++) {
            lattice[i][j] =
                NUM2DBL(rb_ary_entry(rb_ary_entry(r_lattice, i), j));
        }
    }

    num_sym = spgat_get_symmetry(rotations, translations, max_num_sym, lattice,
                                 position, types, num_atom, symprec,
                                 NUM2DBL(r_angle_symprec));

    array = rb_ary_new();

    /* Rotations, translations */
    r_rot = rb_ary_new();
    r_trans = rb_ary_new();
    for (i = 0; i < num_sym; i++) {
        mat = rb_ary_new();
        vec = rb_ary_new();
        for (j = 0; j < 3; j++) {
            rb_ary_push(vec, rb_float_new(translations[i][j]));
            row = rb_ary_new();
            for (k = 0; k < 3; k++) {
                rb_ary_push(row, rb_float_new(rotations[i][j][k]));
            }
            rb_ary_push(mat, row);
        }
        rb_ary_push(r_trans, vec);
        rb_ary_push(r_rot, mat);
    }
    rb_ary_push(array, r_rot);
    rb_ary_push(array, r_trans);

    return array;
}

VALUE method_standardize_cell(VALUE self, VALUE r_lattice, VALUE r_position,
                              VALUE r_types, VALUE r_to_primitive,
                              VALUE r_leave_distrtion, VALUE r_symprec,
                              VALUE r_angle_symprec) {}
