import os
import json
import numpy as np
import random

from itertools import cycle
from skimage import transform
from sklearn.externals import joblib
from sklearn.metrics import log_loss


def binary_crossentropy(y_true, y_pred):
    return log_loss(y_true.flatten(), y_pred.flatten())


def get_bands(band1, band2):
    return np.dstack((
        np.array(band1).reshape(75, 75),
        np.array(band2).reshape(75, 75)
    ))


def get_angle(a):
    if a == 'na':
        return None
    else:
        return a


def get_label(l):
    if l is None:
        return l
    else:
        return int(l)


def unpack(samples):
    images, angles, labels = zip(*[
        (get_bands(s['band_1'], s['band_2']),
        get_angle(s['inc_angle']),
        get_label(s.get('is_iceberg')))
        for s in samples
    ])
    return images, angles, labels


def load_samples(datadir, filename):
    with open(os.path.join(datadir, filename)) as f:
        samples = json.load(f)
    return samples


def train_dev_split(samples, split, shuffle=True):
    split_idx = int(split*len(samples))
    if shuffle:
        random.seed(0)
        random.shuffle(samples)
    return (samples[0:split_idx], samples[split_idx:])


def write_preds(datadir, filename, samples, y_pred):
    with open(os.path.join(datadir, 'predictions', filename), 'w') as f:
        f.write('id,is_iceberg\n')
        f.write('\n'.join([
            '%s,%.6f' % (samples[i].get('id'), y_pred[i][0])
            for i in range(len(samples))
        ]))
        f.write('\n')


def model_fp(datadir, filename):
    return os.path.join(datadir, 'params', filename)


def save_minmax(datadir, filename, minmax):
    joblib.dump(minmax, os.path.join(datadir, 'params', filename))


def load_minmax(datadir, filename):
    return joblib.load(os.path.join(datadir, 'params', filename))


def get_minmax(samples):
    # Extract global mins/max' for normalization
    images, angles, labels = unpack(samples)
    b1_min = min([im[:, :, 0].min() for im in images])
    b1_max = max([im[:, :, 0].max() for im in images])
    b2_min = min([im[:, :, 1].min() for im in images])
    b2_max = max([im[:, :, 1].max() for im in images])
    a_min = min(angles, key=lambda x: x or 1e9)
    a_max = max(angles, key=lambda x: x or -1e9)
    print("Band 1 min/max: %.2f, %.2f" % (b1_min, b1_max))
    print("Band 2 min/max: %.2f, %.2f" % (b2_min, b2_max))
    print("Angles min/max: %.2f, %.2f" % (a_min, a_max))
    return (b1_min, b1_max, b2_min, b2_max, a_min, a_max)


def base_cnn_generator(samples, minmax, batch_size, verbose=False):
    img_dim = 75
    window = 28
    r = int(window/2)
    padding = int(window/np.sqrt(2))
    if verbose:
        print("window is %d px" % window)
        print("padding is %d px" % padding)

    # Fuck the angles.
    images, _, labels = unpack(samples)
    b1_min, b1_max, b2_min, b2_max, _, _ = minmax

    # Yield batches forever, cycling the samples with augmentation:
    # Random rotation and crop, and random Bernoulli mirroring.
    batch_images = []
    batch_labels = []
    for image, label in cycle(zip(images, labels)):
        # Scale each band independently
        img = np.dstack((
            (image[:, :, 0] - b1_min) / (b1_max - b1_min),
            (image[:, :, 1] - b2_min) / (b2_max - b2_min),
        ))

        # Save the maximum of both bands within this image
        b1_peak = img[:, :, 0].max()
        b2_peak = img[:, :, 1].max()

        # Rotate around random midpoint
        row_mid = random.randint(padding, img_dim-padding-1)
        col_mid = random.randint(padding, img_dim-padding-1)
        if verbose:
            print("\nSample %d" % (len(batch_labels)+1))
            print("row/col mid: %d/%d" % (row_mid, col_mid))
            print("Mid values: %.2f, %.2f" % (img[row_mid, col_mid, 0], img[row_mid, col_mid, 1]))

        # NOTE: sklearn.transform.rotate behaves incorrectly when using kwargs 'center'
        # Clip to mid with padding before rotating.
        img = img[row_mid-padding:row_mid+padding, col_mid-padding:col_mid+padding, :]
        rot = random.random() * 360 - 180
        img = transform.rotate(img, rot, order=1)
        if verbose:
            print("Rotation: %.2f deg" % rot)
            print("Mid values: %.2f, %.2f" % (img[padding, padding, 0], img[padding, padding, 1]))
        img = img[padding-r:padding+r, padding-r:padding+r, :]

        # Mirror
        if random.randint(0, 1) and False:
            img = np.fliplr(img)
            if verbose:
                print("Sample was mirrored")

        # Label by checking that subimage contains maximum of either band
        # If the subimage contains a max value that is within 95% of the image
        # max, consider it to contain the feature.
        sub_b1_peak = img[:, :, 0].max()
        sub_b2_peak = img[:, :, 1].max()
        l = int(sub_b1_peak/b1_peak > 0.95 and
                sub_b2_peak/b2_peak > 0.95) * (label + 1)
        if verbose:
            print("Label is %d" % l)
        categorical_label = tuple(int(i==l) for i in range(3))

        batch_images.append(img)
        batch_labels.append(categorical_label)
        if len(batch_labels) == batch_size:
            yield (
                np.array(batch_images, dtype=np.float32),
                np.array(batch_labels, dtype=np.float32)
            )
            batch_images = []
            batch_labels = []


def augment_image(img, mirror, rots):
    if mirror:
        img = np.fliplr(img)
    img = np.rot90(img, k=rots)
    return img


def icenet_generator(samples, minmax, batch_size, crop_offset=3,
                   augment=True, verbose=False):
    img_dim = 75
    mid_x = mid_y = 75.0 / 2
    crop_dim = img_dim - 2*crop_offset
    window = 28
    r = int(window/2)
    
    images, _, labels = unpack(samples)
    b1_min, b1_max, b2_min, b2_max, _, _ = minmax

    # Yield batches forever
    # A batch consists of an input-output tuple
    #   ([X_image_sec1, X_image_sec2, ... , X_image_sec9], Y_labels)
    # where the input image is split over 9 overlapping sections, starting
    # from upper left in row-major order
    batch_image_sections = [[] for _ in range(9)]
    batch_labels = []
    for image, label in cycle(zip(images, labels)):
        # Scale each band independently
        img = np.dstack((
            (image[:, :, 0] - b1_min) / (b1_max - b1_min),
            (image[:, :, 1] - b2_min) / (b2_max - b2_min),
        ))

        # Crop with random offset from midpoint
        row_offset = crop_offset * (2*random.random() - 1)
        col_offset = crop_offset * (2*random.random() - 1)
        if not augment:
            row_offset = col_offset = 0
        row = int(mid_y + row_offset - crop_dim/2)
        col = int(mid_x + col_offset - crop_dim/2)
        img = img[row:row+crop_dim, col:col+crop_dim, :]

        if augment:
            img = augment_image(img, random.randint(0, 1), random.randint(0, 3))

        # Append each section of the image
        for i, ridx in enumerate([r, int(mid_y), crop_dim-r]):
            for j, cidx in enumerate([r, int(mid_x), crop_dim-r]):
                batch_image_sections[i*3+j].append(img[ridx-r:ridx+r, cidx-r:cidx+r, :])

        batch_labels.append(label)
        if len(batch_labels) == batch_size:
            yield (
                [np.array(sec_imgs, dtype=np.float32) for sec_imgs in batch_image_sections],
                np.array(batch_labels, dtype=np.float32)
            )
            batch_image_sections = [[] for _ in range(9)]
            batch_labels = []