import os

import numpy as np
import pandas as pd

from keras.utils import to_categorical
from sklearn.model_selection import KFold, train_test_split


def load_data(path):
    train = pd.read_json(os.path.join(path, "./train.json"))
    test = pd.read_json(os.path.join(path, "./test.json"))
    return (train, test)


def preprocess(df, 
               means=(-22.159262, -24.953745, 40.021883465782651),
               stds=(5.33146, 4.5463958, 4.0815391476694414)):
    X_band_1 = np.array([np.array(band).astype(np.float32).reshape(75, 75) 
                         for band in df["band_1"]])
    X_band_2 = np.array([np.array(band).astype(np.float32).reshape(75, 75)
                         for band in df["band_2"]])

    angl = df['inc_angle'].map(lambda x: np.cos(x * np.pi / 180) if x != 'na' else means[3])
    angl = np.array([np.full(shape=(75, 75), fill_value=angel).astype(np.float32)
                     for angel in angl])

    X_band_1 = (X_band_1 - means[0]) / stds[0]
    X_band_2 = (X_band_2 - means[1]) / stds[1]
    angl = (angl - means[2]) / stds[2]

    images = np.concatenate([X_band_1[:, :, :, np.newaxis],
                             X_band_2[:, :, :, np.newaxis],
                             angl[:, :, :, np.newaxis]],
                            axis=-1)
    return images


def prepare_data_cv(path):
    train, test = load_data(path)
    X_train, y_train = (preprocess(train),
                        to_categorical(train['is_iceberg'].as_matrix().reshape(-1, 1)))

    kfold_data = []
    kf = KFold(n_splits=5, shuffle=True, random_state=0xCAFFE)

    for train_indices, val_indices in kf.split(y_train):
        X_train_cv = X_train[train_indices]
        y_train_cv = y_train[train_indices]

        X_val = X_train[val_indices]
        y_val = y_train[val_indices]

        kfold_data.append((X_train_cv, y_train_cv, X_val, y_val))

    X_test = preprocess(test)

    return (kfold_data, X_test)


def prepare_data(path):
    train, test = load_data(path)
    X_train, y_train = (preprocess(train),
                        to_categorical(train['is_iceberg'].as_matrix().reshape(-1, 1)))

    X_train_cv, X_valid, y_train_cv, y_valid = train_test_split(X_train,
                                                                y_train,
                                                                random_state=0xCAFFE,
                                                                train_size=0.8)

    X_test = preprocess(test)

    return ([(X_train_cv, y_train_cv, X_valid, y_valid)], X_test)

[(32, 'numpy.concatenate', 'np.concatenate', True, 'import numpy as np\n'), (45, 'sklearn.model_selection.KFold', 'KFold', False, 'from sklearn.model_selection import KFold, train_test_split\n'), (66, 'sklearn.model_selection.train_test_split', 'train_test_split', False, 'from sklearn.model_selection import KFold, train_test_split\n'), (11, 'os.path.join', 'os.path.join', False, 'import os\n'), (12, 'os.path.join', 'os.path.join', False, 'import os\n'), (24, 'numpy.cos', 'np.cos', True, 'import numpy as np\n'), (25, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (19, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (21, 'numpy.array', 'np.array', True, 'import numpy as np\n')]