import { assert } from 'chai';
import { logRankTest } from 'pages/resultsView/survival/logRankTest';
import _ from 'lodash';

function getPatientSurvivalCurves(
    exitTime: number[],
    status: number[],
    group: string[]
) {
    return _.values(
        _.groupBy(
            exitTime.map((months, i) => ({
                months,
                status: status[i] === 1,
                group: group[i],
            })),
            'group'
        )
    );
}
describe('logRankTest', () => {
    // validated against R (see logRankGroundTruth.r)
    it('returns the correct p-value when there is no difference in distribution', () => {
        const exitTime = [
            69,
            54,
            75,
            77,
            74,
            59,
            83,
            23,
            19,
            79,
            86,
            96,
            69,
            30,
            71,
            11,
            15,
            57,
            6,
            67,
            72,
            82,
            67,
            98,
            43,
            26,
            3,
            5,
            18,
            51,
            94,
            82,
            47,
            18,
            70,
            67,
            51,
            74,
            3,
            35,
            0,
            11,
            45,
            1,
            83,
            37,
            61,
            58,
            23,
            65,
            4,
            60,
            1,
            98,
            23,
            90,
            80,
            18,
            32,
            59,
            63,
            83,
            75,
            28,
            98,
            73,
            46,
            9,
            27,
            27,
            31,
            61,
            18,
            14,
            77,
            34,
            34,
            40,
            32,
            29,
            71,
            66,
            64,
            44,
            54,
            28,
            5,
            37,
            64,
            69,
            34,
            57,
            9,
            59,
            38,
            32,
            74,
            94,
            95,
            27,
            15,
            86,
            54,
            82,
            88,
            3,
            55,
            60,
            39,
            98,
            89,
            65,
            46,
            69,
            11,
            71,
            98,
            15,
            41,
            62,
            42,
            11,
            3,
            83,
            27,
            37,
            16,
            90,
            77,
            88,
            29,
            62,
            55,
            18,
            81,
            58,
            17,
            69,
            16,
            75,
            48,
            6,
            94,
            90,
            33,
            87,
            46,
            26,
            49,
            67,
            73,
            89,
            72,
            89,
            39,
            67,
            76,
            59,
            77,
            74,
            53,
            80,
            71,
            13,
            85,
            26,
            29,
            98,
            95,
            71,
            43,
            88,
            38,
            54,
            85,
            84,
            67,
            1,
            12,
            47,
            63,
            16,
            72,
            51,
            64,
            9,
            39,
            17,
            95,
            93,
            34,
            95,
            61,
            31,
            79,
            74,
            3,
            42,
            70,
            33,
        ];
        const status = [
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            1,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            1,
            1,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            1,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
        ];
        let group: string[], pValue: number;

        // 2 groups
        group = [
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '1',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '1',
            '2',
            '1',
            '1',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '1',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '1',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '1',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '1',
            '1',
            '1',
            '1',
            '1',
            '1',
            '1',
            '2',
            '1',
            '2',
            '1',
            '1',
            '1',
            '1',
            '2',
            '2',
            '1',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '1',
            '1',
            '1',
            '2',
        ];
        pValue = logRankTest(
            ...getPatientSurvivalCurves(exitTime, status, group)
        )!;
        assert.approximately(pValue, 0.3103677, 0.000001, '2 groups');

        // 5 groups
        group = [
            '3',
            '2',
            '2',
            '2',
            '3',
            '2',
            '4',
            '1',
            '2',
            '5',
            '5',
            '1',
            '4',
            '3',
            '1',
            '1',
            '4',
            '1',
            '3',
            '2',
            '3',
            '1',
            '2',
            '3',
            '3',
            '5',
            '2',
            '5',
            '4',
            '5',
            '4',
            '2',
            '4',
            '2',
            '1',
            '4',
            '1',
            '4',
            '1',
            '4',
            '2',
            '2',
            '4',
            '2',
            '4',
            '4',
            '4',
            '3',
            '1',
            '2',
            '2',
            '4',
            '3',
            '5',
            '1',
            '3',
            '4',
            '5',
            '4',
            '4',
            '3',
            '2',
            '5',
            '2',
            '2',
            '2',
            '1',
            '4',
            '2',
            '4',
            '1',
            '1',
            '1',
            '2',
            '3',
            '4',
            '3',
            '1',
            '1',
            '5',
            '4',
            '3',
            '3',
            '4',
            '3',
            '2',
            '5',
            '3',
            '4',
            '1',
            '5',
            '5',
            '5',
            '3',
            '2',
            '4',
            '4',
            '3',
            '4',
            '4',
            '1',
            '5',
            '1',
            '4',
            '2',
            '3',
            '4',
            '2',
            '5',
            '2',
            '3',
            '3',
            '1',
            '5',
            '1',
            '3',
            '5',
            '5',
            '4',
            '2',
            '1',
            '2',
            '1',
            '2',
            '5',
            '2',
            '2',
            '3',
            '4',
            '5',
            '5',
            '1',
            '2',
            '5',
            '4',
            '5',
            '5',
            '5',
            '2',
            '5',
            '5',
            '1',
            '3',
            '1',
            '3',
            '5',
            '4',
            '1',
            '5',
            '4',
            '3',
            '2',
            '1',
            '2',
            '3',
            '3',
            '3',
            '2',
            '2',
            '5',
            '3',
            '1',
            '3',
            '4',
            '4',
            '4',
            '5',
            '3',
            '4',
            '2',
            '5',
            '1',
            '5',
            '1',
            '2',
            '1',
            '5',
            '2',
            '3',
            '3',
            '1',
            '5',
            '4',
            '1',
            '4',
            '1',
            '1',
            '5',
            '1',
            '1',
            '2',
            '3',
            '5',
            '5',
            '3',
            '5',
            '5',
            '3',
            '3',
            '3',
        ];
        pValue = logRankTest(
            ...getPatientSurvivalCurves(exitTime, status, group)
        )!;
        assert.approximately(pValue, 0.8443513, 0.000001, '5 groups');
    });
    it('returns the correct p-value when distributions are different', () => {
        const exitTime = [
            61,
            5,
            64,
            53,
            92,
            93,
            15,
            45,
            2,
            84,
            54,
            40,
            54,
            86,
            58,
            56,
            51,
            39,
            2,
            78,
            44,
            81,
            32,
            10,
            47,
            37,
            67,
            63,
            87,
            56,
            25,
            90,
            89,
            84,
            25,
            80,
            12,
            29,
            13,
            46,
            36,
            88,
            13,
            17,
            91,
            96,
            87,
            14,
            24,
            14,
            72,
            62,
            62,
            51,
            44,
            2,
            1,
            64,
            57,
            50,
            77,
            10,
            25,
            28,
            76,
            42,
            24,
            7,
            25,
            11,
            50,
            16,
            33,
            22,
            14,
            85,
            94,
            26,
            9,
            66,
            89,
            77,
            95,
            40,
            92,
            15,
            36,
            58,
            52,
            85,
            2,
            50,
            36,
            46,
            47,
            64,
            85,
            50,
            95,
            91,
            22,
            35,
            83,
            3,
            72,
            95,
            78,
            50,
            17,
            90,
            31,
            100,
            87,
            55,
            35,
            72,
            30,
            11,
            74,
            30,
            75,
            85,
            46,
            11,
            24,
            66,
            76,
            13,
            23,
            50,
            98,
            62,
            23,
            99,
            19,
            43,
            53,
            93,
            47,
            57,
            71,
            34,
            57,
            2,
            55,
            71,
            83,
            95,
            16,
            75,
            22,
            35,
            49,
            8,
            19,
            25,
            60,
            68,
            92,
            23,
            25,
            46,
            8,
            81,
            58,
            1,
            96,
            93,
            27,
            92,
            9,
            44,
            1,
            16,
            61,
            5,
            7,
            12,
            86,
            58,
            61,
            66,
            10,
            19,
            62,
            83,
            60,
            35,
            96,
            60,
            97,
            71,
            9,
            44,
            72,
            25,
            80,
            22,
            96,
            15,
        ];
        const status = [
            1,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            1,
            0,
            0,
            1,
            1,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            1,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            1,
            0,
            1,
            1,
            1,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            1,
            1,
            1,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            1,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            0,
            1,
            0,
            1,
            0,
            1,
            0,
            0,
            1,
            0,
            1,
            0,
            0,
            0,
            1,
            1,
            0,
            1,
            1,
            1,
        ];
        let group: string[], pValue: number;

        // 2 groups
        group = [
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '2',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '1',
            '2',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '2',
            '1',
            '1',
            '1',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '2',
            '2',
            '1',
            '2',
            '2',
            '1',
            '1',
            '2',
            '2',
            '1',
            '1',
            '2',
            '1',
            '1',
            '1',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '1',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '2',
            '1',
            '2',
            '2',
            '2',
            '2',
            '1',
            '1',
            '1',
            '2',
            '1',
            '1',
            '2',
            '1',
            '1',
            '1',
            '1',
            '1',
            '2',
            '2',
            '2',
            '2',
            '2',
            '1',
            '2',
        ];
        pValue = logRankTest(
            ...getPatientSurvivalCurves(exitTime, status, group)
        )!;
        assert.approximately(pValue, 0.003960483, 0.000001, '2 groups');

        group = [
            '5',
            '2',
            '5',
            '4',
            '1',
            '1',
            '2',
            '4',
            '2',
            '1',
            '4',
            '4',
            '3',
            '1',
            '4',
            '3',
            '4',
            '3',
            '2',
            '5',
            '4',
            '1',
            '3',
            '2',
            '3',
            '3',
            '5',
            '5',
            '1',
            '4',
            '3',
            '1',
            '5',
            '1',
            '3',
            '5',
            '2',
            '3',
            '2',
            '4',
            '3',
            '1',
            '2',
            '2',
            '1',
            '5',
            '1',
            '2',
            '3',
            '2',
            '5',
            '4',
            '5',
            '4',
            '4',
            '2',
            '2',
            '5',
            '4',
            '4',
            '5',
            '2',
            '3',
            '3',
            '5',
            '4',
            '3',
            '1',
            '3',
            '2',
            '3',
            '2',
            '3',
            '3',
            '2',
            '1',
            '1',
            '3',
            '2',
            '5',
            '1',
            '5',
            '1',
            '4',
            '1',
            '2',
            '3',
            '4',
            '4',
            '1',
            '2',
            '4',
            '3',
            '4',
            '4',
            '5',
            '1',
            '4',
            '1',
            '1',
            '3',
            '3',
            '1',
            '2',
            '5',
            '1',
            '5',
            '3',
            '2',
            '1',
            '3',
            '1',
            '1',
            '4',
            '2',
            '5',
            '3',
            '2',
            '5',
            '3',
            '5',
            '1',
            '4',
            '1',
            '3',
            '5',
            '5',
            '2',
            '3',
            '4',
            '1',
            '5',
            '3',
            '1',
            '2',
            '4',
            '4',
            '1',
            '3',
            '4',
            '5',
            '3',
            '4',
            '2',
            '4',
            '5',
            '1',
            '5',
            '2',
            '5',
            '3',
            '3',
            '4',
            '2',
            '2',
            '3',
            '5',
            '5',
            '1',
            '3',
            '3',
            '4',
            '2',
            '1',
            '4',
            '2',
            '1',
            '1',
            '3',
            '1',
            '2',
            '4',
            '2',
            '2',
            '5',
            '2',
            '2',
            '2',
            '1',
            '4',
            '5',
            '5',
            '2',
            '2',
            '4',
            '1',
            '5',
            '3',
            '5',
            '5',
            '1',
            '5',
            '2',
            '4',
            '5',
            '3',
            '1',
            '3',
            '1',
            '2',
        ];
        pValue = logRankTest(
            ...getPatientSurvivalCurves(exitTime, status, group)
        )!;
        assert.approximately(pValue, 0, 0.000001, '5 groups');
    });
});