# Author: Vlad Niculae <vlad@vene.ro>
# License: BSD 3-clause

import os
import numpy as np
from marseille.argdoc import UkpEssayArgumentationDoc, CdcpArgumentationDoc

cdcp_train_ids = [195, 197, 198, 200, 201, 203, 204, 205, 206, 207, 208, 210,
                  212, 217, 218, 219, 220, 223, 224, 225, 228, 230, 231, 234,
                  239, 302, 305, 306, 308, 309, 312, 314, 315, 318, 321, 324,
                  325, 326, 327, 331, 332, 333, 334, 335, 337, 338, 340, 342,
                  344, 345, 347, 348, 349, 350, 352, 353, 354, 357, 360, 362,
                  364, 370, 371, 373, 375, 377, 379, 380, 381, 382, 383, 385,
                  386, 388, 390, 391, 392, 394, 395, 396, 397, 402, 403, 404,
                  406, 407, 409, 410, 411, 413, 414, 416, 417, 419, 420, 422,
                  423, 424, 426, 427, 428, 430, 433, 435, 436, 439, 442, 444,
                  448, 450, 451, 457, 458, 460, 461, 462, 466, 467, 472, 476,
                  477, 479, 480, 481, 482, 483, 484, 489, 491, 493, 494, 495,
                  496, 498, 499, 500, 501, 502, 503, 505, 506, 512, 513, 514,
                  517, 518, 520, 521, 522, 523, 526, 527, 528, 529, 530, 531,
                  532, 533, 534, 535, 538, 541, 546, 549, 551, 554, 555, 557,
                  559, 560, 563, 564, 566, 567, 568, 571, 572, 574, 575, 576,
                  577, 578, 579, 582, 586, 587, 588, 592, 594, 595, 596, 597,
                  598, 599, 600, 601, 602, 604, 605, 608, 609, 610, 611, 612,
                  614, 621, 622, 623, 626, 627, 628, 631, 633, 635, 639, 644,
                  645, 649, 656, 657, 658, 659, 660, 662, 663, 665, 668, 669,
                  673, 674, 675, 676, 679, 681, 682, 683, 684, 685, 686, 688,
                  689, 691, 692, 694, 695, 703, 704, 705, 710, 717, 720, 721,
                  728, 732, 739, 740, 741, 744, 745, 746, 747, 751, 753, 755,
                  756, 757, 758, 762, 764, 765, 766, 768, 771, 772, 774, 779,
                  783, 784, 785, 786, 788, 793, 794, 797, 798, 802, 803, 805,
                  808, 812, 813, 814, 815, 816, 817, 819, 822, 823, 824, 826,
                  829, 839, 840, 842, 843, 845, 846, 852, 856, 862, 863, 864,
                  865, 867, 870, 872, 881, 882, 883, 884, 888, 889, 891, 901,
                  902, 903, 905, 906, 907, 908, 909, 910, 912, 913, 928, 930,
                  938, 943, 953, 955, 956, 958, 963, 964, 967, 969, 970, 974,
                  975, 977, 980, 987, 988, 989, 990, 993, 994, 995, 996, 997,
                  998, 999, 1000, 1001, 1002, 1006, 1012, 1013, 1014, 1015,
                  1017, 1018, 1020, 1021, 1022, 1024, 1025, 1026, 1027, 1028,
                  1029, 1031, 1032, 1033, 1035, 1036, 1037, 1040, 1041, 1049,
                  1050, 1051, 1053, 1056, 1057, 1059, 1061, 1062, 1063, 1066,
                  1067, 1068, 1069, 1070, 1071, 1072, 1073, 1075, 1076, 1077,
                  1078, 1079, 1083, 1085, 1086, 1088, 1089, 1090, 1091, 1095,
                  1097, 1098, 1101, 1109, 1110, 1111, 1112, 1114, 1118, 1119,
                  1120, 1121, 1125, 1126, 1127, 1128, 1129, 1130, 1131, 1134,
                  1136, 1137, 1140, 1141, 1142, 1144, 1145, 1146, 1147, 1148,
                  1149, 1150, 1154, 1157, 1164, 1165, 1166, 1167, 1168, 1169,
                  1171, 1172, 1173, 1179, 1180, 1181, 1186, 1187, 1189, 1192,
                  1193, 1194, 1196, 1197, 1199, 1201, 1206, 1208, 1209, 1210,
                  1211, 1212, 1217, 1218, 1219, 1226, 1227, 1231, 1234, 1235,
                  1236, 1237, 1238, 1243, 1244, 1246, 1247, 1248, 1249, 1251,
                  1252, 1255, 1256, 1257, 1260, 1261, 1262, 1263, 1265, 1267,
                  1269, 1272, 1273, 1274, 1275, 1280, 1281, 1282, 1284, 1286,
                  1288, 1289, 1292, 1296, 1297, 1298, 1301, 1303, 1307, 1308,
                  1309, 1311, 1312, 1313, 1314, 1316, 1321, 1322, 1323, 1327,
                  1329, 1331, 1337, 1340, 1345, 1347, 1349, 1350, 1353, 1355,
                  1356, 1365, 1368, 1369, 1371, 1373, 1374, 1375, 1377, 1378,
                  1380, 1383, 1384, 1385, 1386, 1387, 1389, 1390, 1394, 1395,
                  1399, 1400, 1401, 1403, 1407, 1408, 1409, 1410, 1412, 1414,
                  1417]

cdcp_test_ids = [194, 196, 199, 226, 236, 238, 240, 303, 310, 320, 336, 339,
                 341, 343, 346, 351, 358, 359, 361, 363, 365, 369, 372, 378,
                 387, 389, 393, 398, 400, 405, 412, 429, 434, 452, 456, 463,
                 464, 485, 486, 487, 488, 492, 524, 525, 536, 540, 543, 562,
                 565, 570, 573, 581, 584, 591, 615, 637, 638, 655, 661, 672,
                 677, 680, 698, 702, 708, 716, 722, 730, 748, 750, 759, 760,
                 763, 776, 790, 796, 799, 806, 810, 811, 821, 825, 835, 837,
                 851, 854, 860, 861, 869, 886, 890, 929, 932, 947, 966, 978,
                 984, 986, 1005, 1016, 1019, 1030, 1044, 1084, 1087, 1092,
                 1093, 1094, 1099, 1108, 1113, 1122, 1123, 1132, 1135, 1151,
                 1152, 1158, 1159, 1178, 1182, 1190, 1191, 1207, 1213, 1220,
                 1239, 1241, 1242, 1250, 1283, 1290, 1295, 1299, 1300, 1318,
                 1324, 1326, 1330, 1332, 1339, 1342, 1344, 1361, 1382, 1388,
                 1398, 1405, 1411, 1418]

ukp_train_ids = [1, 2, 3, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
                 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
                 37, 38, 39, 40, 41, 43, 44, 45, 46, 47, 48, 49, 50, 51, 53,
                 54, 55, 56, 57, 58, 59, 60, 62, 63, 64, 65, 66, 67, 69, 70,
                 73, 74, 75, 76, 78, 79, 80, 81, 83, 84, 85, 87, 88, 89, 90,
                 92, 93, 94, 95, 96, 99, 100, 101, 102, 105, 106, 107, 109,
                 110, 111, 112, 113, 114, 115, 116, 118, 120, 121, 122, 123,
                 124, 125, 127, 128, 130, 131, 132, 133, 134, 135, 137, 138,
                 140, 141, 143, 144, 145, 146, 147, 148, 150, 151, 152, 153,
                 155, 156, 157, 158, 159, 161, 162, 164, 165, 166, 167, 168,
                 170, 171, 173, 174, 175, 176, 177, 178, 179, 181, 183, 184,
                 185, 186, 188, 189, 190, 191, 194, 195, 196, 197, 198, 200,
                 201, 203, 205, 206, 207, 208, 209, 210, 213, 214, 215, 216,
                 217, 219, 222, 223, 224, 225, 226, 228, 230, 231, 232, 233,
                 235, 236, 237, 238, 239, 242, 244, 246, 247, 248, 249, 250,
                 251, 253, 254, 256, 257, 258, 260, 261, 262, 263, 264, 267,
                 268, 269, 270, 271, 272, 273, 274, 275, 276, 279, 280, 281,
                 282, 283, 284, 285, 286, 288, 290, 291, 292, 293, 294, 295,
                 296, 297, 298, 299, 300, 302, 303, 304, 305, 307, 308, 309,
                 311, 312, 313, 314, 315, 317, 318, 319, 320, 321, 323, 324,
                 325, 326, 327, 329, 330, 332, 333, 334, 336, 337, 338, 339,
                 340, 342, 343, 344, 345, 346, 347, 349, 350, 351, 353, 354,
                 356, 357, 358, 360, 361, 362, 363, 365, 366, 367, 368, 369,
                 370, 371, 372, 374, 375, 376, 377, 378, 379, 380, 381, 383,
                 384, 385, 387, 388, 389, 390, 391, 392, 394, 395, 396, 397,
                 399, 400, 401, 402]

ukp_test_ids = [4, 5, 6, 21, 42, 52, 61, 68, 71, 72, 77, 82, 86, 91, 97, 98,
                103, 104, 108, 117, 119, 126, 129, 136, 139, 142, 149, 154,
                160, 163, 169, 172, 180, 182, 187, 192, 193, 199, 202, 204,
                211, 212, 218, 220, 221, 227, 229, 234, 240, 241, 243, 245,
                252, 255, 259, 265, 266, 277, 278, 287, 289, 301, 306, 310,
                316, 322, 328, 331, 335, 341, 348, 352, 355, 359, 364, 373,
                382, 386, 393, 398]

cdcp_ids = sorted(cdcp_train_ids + cdcp_test_ids)
ukp_ids = sorted(ukp_train_ids + ukp_test_ids)


def get_dataset_loader(dataset, split='train'):

    if dataset == 'ukp':
        _tpl = os.path.join("data", "process", "ukp-essays", "essay{:03d}")
        _load = lambda ks: (UkpEssayArgumentationDoc(_tpl.format(k))
                            for k in ks)
        ids = ukp_train_ids if split == 'train' else ukp_test_ids

    elif dataset == 'cdcp':
        _tpl = os.path.join("data", "process", "erule", "{}", "{:05d}")
        _load = lambda ks: (CdcpArgumentationDoc(_tpl.format(split, k))
                            for k in ks)
        ids = cdcp_train_ids if split == 'train' else cdcp_test_ids

    else:
        raise ValueError("Dataset must be 'cdcp' or 'ukp'.")

    return _load, np.array(ids)


def load_embeds(dataset):
    path = os.path.join('data', '{}-glove.npz'.format(dataset))
    try:
        npz = np.load(path)
        return npz['vocab'], npz['embeds']
    except:
        raise ValueError("Could not find embeddings at {}. "
                         "Run `preprocess.py embeddings [...]`".format(path))
