/* -----------------------------------------------------------------------------
 * Robust variance for Cox proportional hazard survival regression
 * -------------------------------------------------------------------------- */

---------------------------------------------------------------------------
-- Data taken from survival package of R
---------------------------------------------------------------------------
DROP TABLE IF EXISTS bladder;
CREATE TABLE bladder(
    id      INTEGER,
    rx      INTEGER,
    num     INTEGER,
    sz      INTEGER,
    stop    INTEGER,
    event   INTEGER,
    enum    INTEGER);

COPY bladder FROM STDIN;
1	1	1	3	1	0	1
1	1	1	3	1	0	2
1	1	1	3	1	0	3
1	1	1	3	1	0	4
2	1	2	1	4	0	1
2	1	2	1	4	0	2
2	1	2	1	4	0	3
2	1	2	1	4	0	4
3	1	1	1	7	0	1
3	1	1	1	7	0	2
3	1	1	1	7	0	3
3	1	1	1	7	0	4
4	1	5	1	10	0	1
4	1	5	1	10	0	2
4	1	5	1	10	0	3
4	1	5	1	10	0	4
5	1	4	1	6	1	1
5	1	4	1	10	0	2
5	1	4	1	10	0	3
5	1	4	1	10	0	4
6	1	1	1	14	0	1
6	1	1	1	14	0	2
6	1	1	1	14	0	3
6	1	1	1	14	0	4
7	1	1	1	18	0	1
7	1	1	1	18	0	2
7	1	1	1	18	0	3
7	1	1	1	18	0	4
8	1	1	3	5	1	1
8	1	1	3	18	0	2
8	1	1	3	18	0	3
8	1	1	3	18	0	4
9	1	1	1	12	1	1
9	1	1	1	16	1	2
9	1	1	1	18	0	3
9	1	1	1	18	0	4
10	1	3	3	23	0	1
10	1	3	3	23	0	2
10	1	3	3	23	0	3
10	1	3	3	23	0	4
11	1	1	3	10	1	1
11	1	1	3	15	1	2
11	1	1	3	23	0	3
11	1	1	3	23	0	4
12	1	1	1	3	1	1
12	1	1	1	16	1	2
12	1	1	1	23	1	3
12	1	1	1	23	0	4
13	1	3	1	3	1	1
13	1	3	1	9	1	2
13	1	3	1	21	1	3
13	1	3	1	23	0	4
14	1	2	3	7	1	1
14	1	2	3	10	1	2
14	1	2	3	16	1	3
14	1	2	3	24	1	4
15	1	1	1	3	1	1
15	1	1	1	15	1	2
15	1	1	1	25	1	3
15	1	1	1	25	0	4
16	1	1	2	26	0	1
16	1	1	2	26	0	2
16	1	1	2	26	0	3
16	1	1	2	26	0	4
17	1	8	1	1	1	1
17	1	8	1	26	0	2
17	1	8	1	26	0	3
17	1	8	1	26	0	4
18	1	1	4	2	1	1
18	1	1	4	26	1	2
18	1	1	4	26	0	3
18	1	1	4	26	0	4
19	1	1	2	25	1	1
19	1	1	2	28	0	2
19	1	1	2	28	0	3
19	1	1	2	28	0	4
20	1	1	4	29	0	1
20	1	1	4	29	0	2
20	1	1	4	29	0	3
20	1	1	4	29	0	4
21	1	1	2	29	0	1
21	1	1	2	29	0	2
21	1	1	2	29	0	3
21	1	1	2	29	0	4
22	1	4	1	29	0	1
22	1	4	1	29	0	2
22	1	4	1	29	0	3
22	1	4	1	29	0	4
23	1	1	6	28	1	1
23	1	1	6	30	1	2
23	1	1	6	30	0	3
23	1	1	6	30	0	4
24	1	1	5	2	1	1
24	1	1	5	17	1	2
24	1	1	5	22	1	3
24	1	1	5	30	0	4
25	1	2	1	3	1	1
25	1	2	1	6	1	2
25	1	2	1	8	1	3
25	1	2	1	12	1	4
26	1	1	3	12	1	1
26	1	1	3	15	1	2
26	1	1	3	24	1	3
26	1	1	3	31	0	4
27	1	1	2	32	0	1
27	1	1	2	32	0	2
27	1	1	2	32	0	3
27	1	1	2	32	0	4
28	1	2	1	34	0	1
28	1	2	1	34	0	2
28	1	2	1	34	0	3
28	1	2	1	34	0	4
29	1	2	1	36	0	1
29	1	2	1	36	0	2
29	1	2	1	36	0	3
29	1	2	1	36	0	4
30	1	3	1	29	1	1
30	1	3	1	36	0	2
30	1	3	1	36	0	3
30	1	3	1	36	0	4
31	1	1	2	37	0	1
31	1	1	2	37	0	2
31	1	1	2	37	0	3
31	1	1	2	37	0	4
32	1	4	1	9	1	1
32	1	4	1	17	1	2
32	1	4	1	22	1	3
32	1	4	1	24	1	4
33	1	5	1	16	1	1
33	1	5	1	19	1	2
33	1	5	1	23	1	3
33	1	5	1	29	1	4
34	1	1	2	41	0	1
34	1	1	2	41	0	2
34	1	1	2	41	0	3
34	1	1	2	41	0	4
35	1	1	1	3	1	1
35	1	1	1	43	0	2
35	1	1	1	43	0	3
35	1	1	1	43	0	4
36	1	2	6	6	1	1
36	1	2	6	43	0	2
36	1	2	6	43	0	3
36	1	2	6	43	0	4
37	1	2	1	3	1	1
37	1	2	1	6	1	2
37	1	2	1	9	1	3
37	1	2	1	44	0	4
38	1	1	1	9	1	1
38	1	1	1	11	1	2
38	1	1	1	20	1	3
38	1	1	1	26	1	4
39	1	1	1	18	1	1
39	1	1	1	48	0	2
39	1	1	1	48	0	3
39	1	1	1	48	0	4
40	1	1	3	49	0	1
40	1	1	3	49	0	2
40	1	1	3	49	0	3
40	1	1	3	49	0	4
41	1	3	1	35	1	1
41	1	3	1	51	0	2
41	1	3	1	51	0	3
41	1	3	1	51	0	4
42	1	1	7	17	1	1
42	1	1	7	53	0	2
42	1	1	7	53	0	3
42	1	1	7	53	0	4
43	1	3	1	3	1	1
43	1	3	1	15	1	2
43	1	3	1	46	1	3
43	1	3	1	51	1	4
44	1	1	1	59	0	1
44	1	1	1	59	0	2
44	1	1	1	59	0	3
44	1	1	1	59	0	4
45	1	3	2	2	1	1
45	1	3	2	15	1	2
45	1	3	2	24	1	3
45	1	3	2	30	1	4
46	1	1	3	5	1	1
46	1	1	3	14	1	2
46	1	1	3	19	1	3
46	1	1	3	27	1	4
47	1	2	3	2	1	1
47	1	2	3	8	1	2
47	1	2	3	12	1	3
47	1	2	3	13	1	4
48	2	1	3	1	0	1
48	2	1	3	1	0	2
48	2	1	3	1	0	3
48	2	1	3	1	0	4
49	2	1	1	1	0	1
49	2	1	1	1	0	2
49	2	1	1	1	0	3
49	2	1	1	1	0	4
50	2	8	1	5	1	1
50	2	8	1	5	0	2
50	2	8	1	5	0	3
50	2	8	1	5	0	4
51	2	1	2	9	0	1
51	2	1	2	9	0	2
51	2	1	2	9	0	3
51	2	1	2	9	0	4
52	2	1	1	10	0	1
52	2	1	1	10	0	2
52	2	1	1	10	0	3
52	2	1	1	10	0	4
53	2	1	1	13	0	1
53	2	1	1	13	0	2
53	2	1	1	13	0	3
53	2	1	1	13	0	4
54	2	2	6	3	1	1
54	2	2	6	14	0	2
54	2	2	6	14	0	3
54	2	2	6	14	0	4
55	2	5	3	1	1	1
55	2	5	3	3	1	2
55	2	5	3	5	1	3
55	2	5	3	7	1	4
56	2	5	1	18	0	1
56	2	5	1	18	0	2
56	2	5	1	18	0	3
56	2	5	1	18	0	4
57	2	1	3	17	1	1
57	2	1	3	18	0	2
57	2	1	3	18	0	3
57	2	1	3	18	0	4
58	2	5	1	2	1	1
58	2	5	1	19	0	2
58	2	5	1	19	0	3
58	2	5	1	19	0	4
59	2	1	1	17	1	1
59	2	1	1	19	1	2
59	2	1	1	21	0	3
59	2	1	1	21	0	4
60	2	1	1	22	0	1
60	2	1	1	22	0	2
60	2	1	1	22	0	3
60	2	1	1	22	0	4
61	2	1	3	25	0	1
61	2	1	3	25	0	2
61	2	1	3	25	0	3
61	2	1	3	25	0	4
62	2	1	5	25	0	1
62	2	1	5	25	0	2
62	2	1	5	25	0	3
62	2	1	5	25	0	4
63	2	1	1	25	0	1
63	2	1	1	25	0	2
63	2	1	1	25	0	3
63	2	1	1	25	0	4
64	2	1	1	6	1	1
64	2	1	1	12	1	2
64	2	1	1	13	1	3
64	2	1	1	26	0	4
65	2	1	1	6	1	1
65	2	1	1	27	0	2
65	2	1	1	27	0	3
65	2	1	1	27	0	4
66	2	2	1	2	1	1
66	2	2	1	29	0	2
66	2	2	1	29	0	3
66	2	2	1	29	0	4
67	2	8	3	26	1	1
67	2	8	3	35	1	2
67	2	8	3	36	0	3
67	2	8	3	36	0	4
68	2	1	1	38	0	1
68	2	1	1	38	0	2
68	2	1	1	38	0	3
68	2	1	1	38	0	4
69	2	1	1	22	1	1
69	2	1	1	23	1	2
69	2	1	1	27	1	3
69	2	1	1	32	1	4
70	2	6	1	4	1	1
70	2	6	1	16	1	2
70	2	6	1	23	1	3
70	2	6	1	27	1	4
71	2	3	1	24	1	1
71	2	3	1	26	1	2
71	2	3	1	29	1	3
71	2	3	1	40	1	4
72	2	3	2	41	0	1
72	2	3	2	41	0	2
72	2	3	2	41	0	3
72	2	3	2	41	0	4
73	2	1	1	41	0	1
73	2	1	1	41	0	2
73	2	1	1	41	0	3
73	2	1	1	41	0	4
74	2	1	1	1	1	1
74	2	1	1	27	1	2
74	2	1	1	43	0	3
74	2	1	1	43	0	4
75	2	1	1	44	0	1
75	2	1	1	44	0	2
75	2	1	1	44	0	3
75	2	1	1	44	0	4
76	2	6	1	2	1	1
76	2	6	1	20	1	2
76	2	6	1	23	1	3
76	2	6	1	27	1	4
77	2	1	2	45	0	1
77	2	1	2	45	0	2
77	2	1	2	45	0	3
77	2	1	2	45	0	4
78	2	1	4	2	1	1
78	2	1	4	46	0	2
78	2	1	4	46	0	3
78	2	1	4	46	0	4
79	2	1	4	46	0	1
79	2	1	4	46	0	2
79	2	1	4	46	0	3
79	2	1	4	46	0	4
80	2	3	3	49	0	1
80	2	3	3	49	0	2
80	2	3	3	49	0	3
80	2	3	3	49	0	4
81	2	1	1	50	0	1
81	2	1	1	50	0	2
81	2	1	1	50	0	3
81	2	1	1	50	0	4
82	2	4	1	4	1	1
82	2	4	1	24	1	2
82	2	4	1	47	1	3
82	2	4	1	50	0	4
83	2	3	4	54	0	1
83	2	3	4	54	0	2
83	2	3	4	54	0	3
83	2	3	4	54	0	4
84	2	2	1	38	1	1
84	2	2	1	54	0	2
84	2	2	1	54	0	3
84	2	2	1	54	0	4
85	2	1	3	59	0	1
85	2	1	3	59	0	2
85	2	1	3	59	0	3
85	2	1	3	59	0	4
\.

---------------------------------------------------------------------------
-- no stratification
---------------------------------------------------------------------------
drop table if exists bladder_coxph_out;
drop table if exists bladder_coxph_out_summary;
select coxph_train('bladder', 'bladder_coxph_out', 'stop', 'ARRAY[rx, sz, num]', 'event=1', NULL, NULL);

select * from bladder_coxph_out;
select * from bladder_coxph_out_summary;

-- robust variance
drop table if exists bladder_rv_coxph_out;
select robust_variance_coxph('bladder_coxph_out', 'bladder_rv_coxph_out');

SELECT assert(
    relative_error(coef, ARRAY[-0.53838728618390019, -0.05439102328789678, 0.19234011205408327]) < 1e-6 AND
    relative_error(loglikelihood, -589.06745420305799144) < 1e-6 AND
    relative_error(std_err, ARRAY[0.19985872786736181, 0.06960069357280596, 0.04646488068044735]) < 1e-6 AND
    relative_error(robust_se, ARRAY[0.19296810565837841, 0.07081459654480815, 0.04766700334445047]) < 1e-6 AND
    relative_error(robust_z, ARRAY[-2.7900324996558528, -0.7680764410410881, 4.0350787454415480]) < 1e-6 AND
    relative_error(robust_p, ARRAY[0.00527027509779776, 0.442441772410999, 0.0000545839882079013]) < 1e-6,
    'Cox Robust Variance Estimators (bladder w/o stratification): Wrong results'
) FROM bladder_rv_coxph_out;

-- cluster variance (id)
drop table if exists bladder_cv_coxph_out;
select clustered_variance_coxph('bladder_coxph_out', 'bladder_cv_coxph_out', 'id');

SELECT assert(
    relative_error(coef, ARRAY[-0.53838728618390019, -0.05439102328789678, 0.19234011205408327]) < 1e-6 AND
    relative_error(loglikelihood, -589.06745420305799144) < 1e-6 AND
    relative_error(std_err, ARRAY[0.19985872786736181, 0.06960069357280596, 0.04646488068044735]) < 1e-6 AND
    relative_error(clustered_se, ARRAY[0.29467251572845515, 0.08753417995257935, 0.06689308222921776]) < 1e-6 AND
    relative_error(clustered_z, ARRAY[-1.8270699079381791, -0.6213689705822628, 2.8753363672943806]) < 1e-6 AND
    relative_error(clustered_p, ARRAY[0.06768925905530210, 0.53435688328289499, 0.00403597282368529]) < 1e-6,
    'Cox Cluster Variance Estimators (bladder w/ cluster(id) w/o stratification): Wrong results'
) FROM bladder_cv_coxph_out;

-- cluster variance (id, enum)
drop table if exists bladder_cv_coxph_out;
select clustered_variance_coxph('bladder_coxph_out', 'bladder_cv_coxph_out', 'id, enum');

SELECT assert(
    relative_error(coef, ARRAY[-0.53838728618390019, -0.05439102328789678, 0.19234011205408327]) < 1e-6 AND
    relative_error(loglikelihood, -589.06745420305799144) < 1e-6 AND
    relative_error(std_err, ARRAY[0.19985872786736181, 0.06960069357280596, 0.04646488068044735]) < 1e-6 AND
    relative_error(clustered_se, ARRAY[0.19296810565837841, 0.07081459654480815, 0.04766700334445047]) < 1e-6 AND
    relative_error(clustered_z, ARRAY[-2.7900324996558528, -0.7680764410410881, 4.0350787454415480]) < 1e-6 AND
    relative_error(clustered_p, ARRAY[0.00527027509779776, 0.442441772410999, 0.0000545839882079013]) < 1e-6,
    'Cox Cluster Variance Estimators (bladder w/ cluster(id) + cluster(enum) w/o stratification): Wrong results'
) FROM bladder_cv_coxph_out;


---------------------------------------------------------------------------
-- with stratification
---------------------------------------------------------------------------
drop table if exists bladder_coxph_out;
drop table if exists bladder_coxph_out_summary;
select coxph_train('bladder', 'bladder_coxph_out', 'stop', 'ARRAY[rx, sz, num]', 'event=1', 'enum', NULL);

-- robust variance
drop table if exists bladder_rv_coxph_out;
select robust_variance_coxph('bladder_coxph_out', 'bladder_rv_coxph_out');

SELECT assert(
    relative_error(coef, ARRAY[-0.57986081953852453, -0.05093865053703769, 0.20849139987047763]) < 1e-6 AND
    relative_error(loglikelihood, -428.05790500373399254) < 1e-6 AND
    relative_error(std_err, ARRAY[0.20117781307700300, 0.06966981665067325, 0.04691285708902546]) < 1e-6 AND
    relative_error(robust_se, ARRAY[0.18929312730428571, 0.06364037833687450, 0.04464477664866482]) < 1e-6 AND
    relative_error(robust_z, ARRAY[-3.0632956821850557, -0.8004140117992169, 4.6700065611530599]) < 1e-6 AND
    relative_error(robust_p, ARRAY[0.00218913638636697, 0.423470965451148, 0.00000301190115603589]) < 1e-6,
    'Cox Robust Variance Estimators (bladder w/ stratification): Wrong results'
) FROM bladder_rv_coxph_out;

-- cluster variance
drop table if exists bladder_cv_coxph_out;
select clustered_variance_coxph('bladder_coxph_out', 'bladder_cv_coxph_out', 'id');

SELECT assert(
    relative_error(coef, ARRAY[-0.57986081953852453, -0.05093865053703769, 0.20849139987047763]) < 1e-6 AND
    relative_error(loglikelihood, -428.05790500373399254) < 1e-6 AND
    relative_error(std_err, ARRAY[0.20117781307700300, 0.06966981665067325, 0.04691285708902546]) < 1e-6 AND
    relative_error(clustered_se, ARRAY[0.30343529013002124, 0.09303573931640849, 0.06567477635846880]) < 1e-6 AND
    relative_error(clustered_z, ARRAY[-1.9109867520355186, -0.5475170177752731, 3.1746038803768615]) < 1e-6 AND
    relative_error(clustered_p, ARRAY[0.05600628340725960, 0.58402358222142303, 0.00150041173830195]) < 1e-6,
    'Cox Cluster Variance Estimators (bladder w/ cluster(id) w/ stratification): Wrong results'
) FROM bladder_cv_coxph_out;