// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

import Arrow, { RecordBatch } from '../Arrow';

const { predicate, Table } = Arrow;

const { col, lit, custom } = predicate;

const F32 = 0, I32 = 1, DICT = 2;
const test_data = [
    {name: `single record batch`,
     table: Table.from({
          'schema': {
            'fields': [
              {
                'name': 'f32',
                'type': {
                  'name': 'floatingpoint',
                  'precision': 'SINGLE'
                },
                'nullable': false,
                'children': [],
              },
              {
                'name': 'i32',
                'type': {
                  'name': 'int',
                  'isSigned': true,
                  'bitWidth': 32
                },
                'nullable': false,
                'children': [],
              },
              {
                'name': 'dictionary',
                'type': {
                  'name': 'utf8'
                },
                'nullable': false,
                'children': [],
                'dictionary': {
                  'id': 0,
                  'indexType': {
                    'name': 'int',
                    'isSigned': true,
                    'bitWidth': 8
                  },
                  'isOrdered': false
                }
              }
            ]
          },
          'dictionaries': [{
            'id': 0,
            'data': {
              'count': 3,
              'columns': [
                {
                  'name': 'DICT0',
                  'count': 3,
                  'VALIDITY': [],
                  'OFFSET': [
                    0,
                    1,
                    2,
                    3
                  ],
                  'DATA': [
                    'a',
                    'b',
                    'c',
                  ]
                }
              ]
            }
          }],
          'batches': [{
            'count': 7,
            'columns': [
              {
                'name': 'f32',
                'count': 7,
                'VALIDITY': [],
                'DATA': [-0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3]
              },
              {
                'name': 'i32',
                'count': 7,
                'VALIDITY': [],
                'DATA': [-1, 1, -1, 1, -1, 1, -1]
              },
              {
                'name': 'dictionary',
                'count': 7,
                'VALIDITY': [],
                'DATA': [0, 1, 2, 0, 1, 2, 0]
              }
            ]
          }]
        }),
        // Use Math.fround to coerce to float32
     values: [
         [Math.fround(-0.3), -1, 'a'],
         [Math.fround(-0.2),  1, 'b'],
         [Math.fround(-0.1), -1, 'c'],
         [Math.fround( 0  ),  1, 'a'],
         [Math.fround( 0.1), -1, 'b'],
         [Math.fround( 0.2),  1, 'c'],
         [Math.fround( 0.3), -1, 'a']
     ]},
     {name: `multiple record batches`,
      table: Table.from({
          'schema': {
            'fields': [
              {
                'name': 'f32',
                'type': {
                  'name': 'floatingpoint',
                  'precision': 'SINGLE'
                },
                'nullable': false,
                'children': [],
              },
              {
                'name': 'i32',
                'type': {
                  'name': 'int',
                  'isSigned': true,
                  'bitWidth': 32
                },
                'nullable': false,
                'children': [],
              },
              {
                'name': 'dictionary',
                'type': {
                  'name': 'utf8'
                },
                'nullable': false,
                'children': [],
                'dictionary': {
                  'id': 0,
                  'indexType': {
                    'name': 'int',
                    'isSigned': true,
                    'bitWidth': 8
                  },
                  'isOrdered': false
                }
              }
            ]
          },
          'dictionaries': [{
            'id': 0,
            'data': {
              'count': 3,
              'columns': [
                {
                  'name': 'DICT0',
                  'count': 3,
                  'VALIDITY': [],
                  'OFFSET': [
                    0,
                    1,
                    2,
                    3
                  ],
                  'DATA': [
                    'a',
                    'b',
                    'c',
                  ]
                }
              ]
            }
          }],
          'batches': [{
            'count': 3,
            'columns': [
              {
                'name': 'f32',
                'count': 3,
                'VALIDITY': [],
                'DATA': [-0.3, -0.2, -0.1]
              },
              {
                'name': 'i32',
                'count': 3,
                'VALIDITY': [],
                'DATA': [-1, 1, -1]
              },
              {
                'name': 'dictionary',
                'count': 3,
                'VALIDITY': [],
                'DATA': [0, 1, 2]
              }
            ]
          }, {
            'count': 3,
            'columns': [
              {
                'name': 'f32',
                'count': 3,
                'VALIDITY': [],
                'DATA': [0, 0.1, 0.2]
              },
              {
                'name': 'i32',
                'count': 3,
                'VALIDITY': [],
                'DATA': [1, -1, 1]
              },
              {
                'name': 'dictionary',
                'count': 3,
                'VALIDITY': [],
                'DATA': [0, 1, 2]
              }
            ]
          }, {
            'count': 3,
            'columns': [
              {
                'name': 'f32',
                'count': 3,
                'VALIDITY': [],
                'DATA': [0.3, 0.2, 0.1]
              },
              {
                'name': 'i32',
                'count': 3,
                'VALIDITY': [],
                'DATA': [-1, 1, -1]
              },
              {
                'name': 'dictionary',
                'count': 3,
                'VALIDITY': [],
                'DATA': [0, 1, 2]
              }
            ]
          }]
      }),
      values: [
            [Math.fround(-0.3), -1, 'a'],
            [Math.fround(-0.2),  1, 'b'],
            [Math.fround(-0.1), -1, 'c'],
            [Math.fround( 0  ),  1, 'a'],
            [Math.fround( 0.1), -1, 'b'],
            [Math.fround( 0.2),  1, 'c'],
            [Math.fround( 0.3), -1, 'a'],
            [Math.fround( 0.2),  1, 'b'],
            [Math.fround( 0.1), -1, 'c'],
      ]}
];

describe(`Table`, () => {
    test(`can create an empty table`, () => {
        expect(Table.empty().length).toEqual(0);
    });
    test(`Table.from([]) creates an empty table`, () => {
        expect(Table.from([]).length).toEqual(0);
    });
    test(`Table.from() creates an empty table`, () => {
        expect(Table.from().length).toEqual(0);
    });
    for (let datum of test_data) {
        describe(datum.name, () => {
            const table = datum.table;
            const values = datum.values;

            test(`has the correct length`, () => {
                expect(table.length).toEqual(values.length);
            });
            test(`gets expected values`, () => {
                for (let i = -1; ++i < values.length;) {
                    expect(table.get(i).toArray()).toEqual(values[i]);
                }
            });
            test(`iterates expected values`, () => {
                let i = 0;
                for (let row of table) {
                    expect(row.toArray()).toEqual(values[i++]);
                }
            });
            describe(`scan()`, () => {
                test(`yields all values`, () => {
                    let expected_idx = 0;
                    table.scan((idx, batch) => {
                        const columns = batch.schema.fields.map((_, i) => batch.getChildAt(i)!);
                        expect(columns.map((c) => c.get(idx))).toEqual(values[expected_idx++]);
                    });
                });
                test(`calls bind function with every batch`, () => {
                    let bind = jest.fn();
                    table.scan(() => {}, bind);
                    for (let batch of table.batches) {
                        expect(bind).toHaveBeenCalledWith(batch);
                    }
                });
            });
            test(`count() returns the correct length`, () => {
                expect(table.count()).toEqual(values.length);
            });
            test(`getColumnIndex`, () => {
                expect(table.getColumnIndex('i32')).toEqual(I32);
                expect(table.getColumnIndex('f32')).toEqual(F32);
                expect(table.getColumnIndex('dictionary')).toEqual(DICT);
            });
            let get_i32: (idx: number) => number, get_f32: (idx: number) => number;
            const filter_tests = [
                {
                    name:     `filter on f32 >= 0`,
                    filtered: table.filter(col('f32').gteq(0)),
                    expected: values.filter((row) => row[F32] >= 0)
                }, {
                    name:     `filter on 0 <= f32`,
                    filtered: table.filter(lit(0).lteq(col('f32'))),
                    expected: values.filter((row) => 0 <= row[F32])
                }, {
                    name:     `filter on i32 <= 0`,
                    filtered: table.filter(col('i32').lteq(0)),
                    expected: values.filter((row) => row[I32] <= 0)
                }, {
                    name:     `filter on 0 >= i32`,
                    filtered: table.filter(lit(0).gteq(col('i32'))),
                    expected: values.filter((row) => 0 >= row[I32])
                }, {
                    name:     `filter on f32 <= -.25 || f3 >= .25`,
                    filtered: table.filter(col('f32').lteq(-.25).or(col('f32').gteq(.25))),
                    expected: values.filter((row) => row[F32] <= -.25 || row[F32] >= .25)
                }, {
                    name:     `filter method combines predicates (f32 >= 0 && i32 <= 0)`,
                    filtered: table.filter(col('i32').lteq(0)).filter(col('f32').gteq(0)),
                    expected: values.filter((row) => row[I32] <= 0 && row[F32] >= 0)
                }, {
                    name:     `filter on dictionary == 'a'`,
                    filtered: table.filter(col('dictionary').eq('a')),
                    expected: values.filter((row) => row[DICT] === 'a')
                }, {
                    name:     `filter on 'a' == dictionary (commutativity)`,
                    filtered: table.filter(lit('a').eq(col('dictionary'))),
                    expected: values.filter((row) => row[DICT] === 'a')
                }, {
                    name:     `filter on f32 >= i32`,
                    filtered: table.filter(col('f32').gteq(col('i32'))),
                    expected: values.filter((row) => row[F32] >= row[I32])
                }, {
                    name:     `filter on f32 <= i32`,
                    filtered: table.filter(col('f32').lteq(col('i32'))),
                    expected: values.filter((row) => row[F32] <= row[I32])
                }, {
                    name:     `filter on f32*i32 > 0 (custom predicate)`,
                    filtered: table.filter(custom(
                        (idx: number) => (get_f32(idx) * get_i32(idx) > 0),
                        (batch: RecordBatch) => {
                            get_f32 = col('f32').bind(batch);
                            get_i32 = col('i32').bind(batch);
                        })),
                    expected: values.filter((row) => (row[F32] as number) * (row[I32] as number) > 0)
                }
            ];
            for (let this_test of filter_tests) {
                const { name, filtered, expected } = this_test;
                describe(name, () => {
                    test(`count() returns the correct length`, () => {
                        expect(filtered.count()).toEqual(expected.length);
                    });
                    describe(`scan()`, () => {
                        test(`iterates over expected values`, () => {
                            let expected_idx = 0;
                            filtered.scan((idx, batch) => {
                                const columns = batch.schema.fields.map((_, i) => batch.getChildAt(i)!);
                                expect(columns.map((c) => c.get(idx))).toEqual(expected[expected_idx++]);
                            });
                        });
                        test(`calls bind function on every batch`, () => {
                            // Techincally, we only need to call bind on
                            // batches with data that match the predicate, so
                            // this test may fail in the future if we change
                            // that - and that's ok!
                            let bind = jest.fn();
                            filtered.scan(() => {}, bind);
                            for (let batch of table.batches) {
                                expect(bind).toHaveBeenCalledWith(batch);
                            }
                        });
                    });
                });
            }
            test(`countBy on dictionary returns the correct counts`, () => {
                // Make sure countBy works both with and without the Col wrapper
                // class
                let expected: {[key: string]: number} = {'a': 0, 'b': 0, 'c': 0};
                for (let row of values) {
                    expected[row[DICT]] += 1;
                }

                expect(table.countBy(col('dictionary')).toJSON()).toEqual(expected);
                expect(table.countBy('dictionary').toJSON()).toEqual(expected);
            });
            test(`countBy on dictionary with filter returns the correct counts`, () => {
                let expected: {[key: string]: number} = {'a': 0, 'b': 0, 'c': 0};
                for (let row of values) {
                    if (row[I32] === 1) { expected[row[DICT]] += 1; }
                }

                expect(table.filter(col('i32').eq(1)).countBy('dictionary').toJSON()).toEqual(expected);
            });
            test(`countBy on non dictionary column throws error`, () => {
                expect(() => { table.countBy('i32'); }).toThrow();
                expect(() => { table.filter(col('dict').eq('a')).countBy('i32'); }).toThrow();
            });
            test(`countBy on non-existent column throws error`, () => {
                expect(() => { table.countBy('FAKE'); }).toThrow();
            });
            test(`table.select() basic tests`, () => {
                let selected = table.select('f32', 'dictionary');
                expect(selected.schema.fields.length).toEqual(2);
                expect(selected.schema.fields[0]).toEqual(table.schema.fields[0]);
                expect(selected.schema.fields[1]).toEqual(table.schema.fields[2]);

                expect(selected.length).toEqual(values.length);
                let idx = 0, expected_row;
                for (let row of selected) {
                    expected_row = values[idx++];
                    expect(row.get(0)).toEqual(expected_row[F32]);
                    expect(row.get(1)).toEqual(expected_row[DICT]);
                }
            });
            test(`table.toString()`, () => {
                let selected = table.select('i32', 'dictionary');
                let headers  = [`"row_id"`, `"i32: Int32"`, `"dictionary: Dictionary<Int8, Utf8>"`];
                let expected = [headers.join(' | '), ...values.map((row, idx) => {
                    return [`${idx}`, `${row[I32]}`, `"${row[DICT]}"`].map((str, col) => {
                                return leftPad(str, ' ', headers[col].length);
                            }).join(' | ');
                })].join('\n') + '\n';
                expect(selected.toString()).toEqual(expected);
            });
            test(`table.filter(..).count() on always false predicates returns 0`, () => {
                expect(table.filter(col('i32').gteq(100)).count()).toEqual(0);
                expect(table.filter(col('dictionary').eq('z')).count()).toEqual(0);
            });
            describe(`lit-lit comparison`, () => {
                test(`always-false count() returns 0`, () => {
                    expect(table.filter(lit('abc').eq('def')).count()).toEqual(0);
                    expect(table.filter(lit(0).gteq(1)).count()).toEqual(0);
                });
                test(`always-true count() returns length`, () => {
                    expect(table.filter(lit('abc').eq('abc')).count()).toEqual(table.length);
                    expect(table.filter(lit(-100).lteq(0)).count()).toEqual(table.length);
                });
            });
            describe(`col-col comparison`, () => {
                test(`always-false count() returns 0`, () => {
                    expect(table.filter(col('dictionary').eq(col('i32'))).count()).toEqual(0);
                });
                test(`always-true count() returns length`, () => {
                    expect(table.filter(col('dictionary').eq(col('dictionary'))).count()).toEqual(table.length);
                });
            });
        });
    }
});

function leftPad(str: string, fill: string, n: number) {
    return (new Array(n + 1).join(fill) + str).slice(-1 * n);
}