# -*- coding: utf-8 -*-
# @Author: Aliaksei Halachkin
# @Date:   2016-10-23 12:29:56
# @Last Modified by:   Aliaksei Halachkin
# @Last Modified time: 2017-04-22 19:04:20

import json
import pandas as pd
from sklearn.cross_validation import StratifiedKFold

dataset = "..\\..\\training_data\\trainset_v3_profile_clusters.csv"
save_to = "..\\configs\\folds.json"
n_folds = 5

with open('..\\configs\\class_n_clusters.json', 'r') as f:
    clusters_conf = json.load(f)

df = pd.read_csv(dataset, sep=';')
exclude_classes = clusters_conf['exclude_classes']

df = df[~df.class_name.isin(exclude_classes)]

df['class_name_pcluster'] = df.class_name + ' PCluster: ' + df.PCluster.map(str)

y = df['class_name_pcluster'].as_matrix()

skf = StratifiedKFold(y, n_folds=n_folds, shuffle=True, random_state=0)

folds = []
for it, (train_ind, test_ind) in enumerate(skf):
    folds.append( { "id": it,
                    "train_ind": list(train_ind), 
                    "test_ind" : list(test_ind)})


info = dict(dataset=dataset, folds=folds)

json.dump(info, open(save_to, 'w'), indent=4, sort_keys=True)