import os

import numpy as np

import data_transforms

import pathfinder

import utils

import utils_lung

from configuration import set_configuration, config

from utils_plots import plot_2d_animation, plot_slice_3d_3

set_configuration('test_config')

def test_luna3d():

    # path = '/mnt/sda3/data/kaggle-lung/lunapred/luna_scan_v3_dice-20170131-173443/'

    path = '/mnt/sda3/data/kaggle-lung/lunapred_el/luna_scan_v3_dice-20170201-231707/'

    files = os.listdir(path)

    print files

    x, y, p = [], [], []

    for f in files:

        if 'in' in f:

            x.append(f)

        elif 'tgt' in f:

            y.append(f)

        else:

            p.append(f)

    x = sorted(x)

    y = sorted(y)

    p = sorted(p)

    for xf, yf, pf in zip(x, y, p):

        x_batch = utils.load_pkl(path + xf)

        pred_batch = utils.load_pkl(path + pf)

        y_batch = utils.load_pkl(path + yf)

        print xf

        print yf

        print pf

        # plot_2d_animation(x_batch[0], y_batch[0], pred_batch[0])

        plot_slice_3d_3(x_batch[0,0],y_batch[0,0],pred_batch[0,0],0,'aa')

if __name__ == '__main__':

    test_luna3d()