import sys
import numpy as np
import cv2

import pickle
import cortex.utils
import cortex.vision

import colormap

from frame_interp import decompose
import cupy as cp

# Initialize CuPy
xp = cp
print('Using GPU.')
xp.cuda.Device(0).use()

def chunks(l, n):
    for i in range(0, len(l), n):
        yield l[i:i + n]

def pyr_to_cpu(pyr):

    pyr_cpu = {'pind': pyr['pind'], 'pyramids': list(),
               'amplitude': list(), 'phase': list()}

    for key in ['amplitude', 'phase']:
        for channel_idx in range(len(pyr[key])):
            pyr_channel = pyr[key][channel_idx]
            pyr_channel_cpu = list()
            for pyr_level_band in pyr_channel:
                x_cpu = xp.asnumpy(pyr_level_band).astype(np.float32)
                pyr_channel_cpu.append(x_cpu)

            pyr_cpu[key].append(pyr_channel_cpu)

    for key in ['high_pass', 'low_pass']:
        pyr_cpu[key] = xp.asnumpy(key)

    for key in pyr.keys():
        if key == 'pyramids': continue
        if key in pyr_cpu.keys(): continue
        print('copying {}'.format(key))
        pyr_cpu[key] = pyr[key]

    return pyr_cpu

def extract_phase(image, num_levels=5, num_orientations=8, t_width=1, scale=0.5):

    height, width = image.shape[0], image.shape[1]
    num_channels = 1 if image.ndim == 2 else image.shape[2]

    # Perform the actual decomposition on the GPU
    # ['low_pass', 'pind', 'amplitude', 'pyramids', 'phase', 'high_pass']
    pyr = decompose(image, num_levels, num_orientations, t_width, scale, num_levels, xp)

    pyr['orig_height'] = height
    pyr['orig_width'] = width
    pyr['orig_channels'] = num_channels
    pyr['num_levels'] = num_levels
    pyr['num_orientations'] = num_orientations

    pyr['pyr_height']  = pyr['pyramids'][0].ht
    pyr['pyr_order']   = pyr['pyramids'][0].order
    pyr['pyr_twidth']  = pyr['pyramids'][0].twidth
    pyr['pyr_scale']   = pyr['pyramids'][0].scale
    pyr['pyr_nscales'] = pyr['pyramids'][0].n_scales

    return pyr

def visualize_phase(pyr, verbose=True):

    for channel_idx in range(pyr['orig_channels']):

        # Split pyramid into the levels
        pyr_levels = list(chunks(pyr['phase'][channel_idx], pyr['num_orientations']))
        num_levels = len(pyr_levels)

        for level_idx in range(pyr['num_levels']):

            print('#'*60)
            pyr_orientations = pyr_levels[level_idx]
            num_orientations = len(pyr_orientations)

            bands_sum = xp.zeros_like(pyr_orientations[0])

            for orientation_idx in range(num_orientations):

                pyr_band = pyr_orientations[orientation_idx]
                bands_sum += pyr_band

                if verbose:
                    print('Image {}x{} | Channel {} | Level {} | Orientation {} | {}'.format(
                        pyr['orig_height'], pyr['orig_width'], channel_idx, level_idx, orientation_idx, pyr_band.shape))

            # Visualization
            pyr_band_viz = bands_sum.get().astype(np.float32)  # move back to CPU
            pyr_band_viz = cortex.vision.normalize_for_display(pyr_band_viz)
            #pyr_band_viz = (pyr_band_viz*255).astype(np.uint8)
            #pyr_band_viz = colormap.apply_custom_colormap(pyr_band_viz)

            cv2.imshow('level {}'.format(level_idx), pyr_band_viz)

    cv2.waitKey(10)
    return pyr

##########################################################################################
##########################################################################################

main_dir = '/home/tomrunia/data/WindDataset/videos_as_images/'
video_subdirs = cortex.utils.find_subdirs(main_dir)

cnt = 0

for video_subdir in video_subdirs:
    image_files = cortex.utils.find_files(video_subdir, extensions='jpg')
    cnt += len(image_files)

print(cnt)
exit()

for image_file in image_files:
    image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)
    pyr = extract_phase(image, num_orientations=8, t_width=1)

    # Move the entire thing back to the CPU
    pyr_cpu = pyr_to_cpu(pyr)


    pickle.dump(pyr_cpu, open('tmp.pkl', 'w'), protocol=-1)
    exit()



