"""Auxiliary methods."""
import os
import json
from errno import EEXIST

import numpy as np
import seaborn as sns
import cPickle as pickle
import matplotlib.pyplot as plt

sns.set()

DEFAULT_LOG_DIR = 'log'
ATOB_WEIGHTS_FILE = 'atob_weights.h5'
D_WEIGHTS_FILE = 'd_weights.h5'


class MyDict(dict):
    """
    Dictionary that allows to access elements with dot notation.

    ex:
        >> d = MyDict({'key': 'val'})
        >> d.key
        'val'
        >> d.key2 = 'val2'
        >> d
        {'key2': 'val2', 'key': 'val'}
    """

    __getattr__ = dict.get
    __setattr__ = dict.__setitem__


def convert_to_rgb(img, is_binary=False):
    """Given an image, make sure it has 3 channels and that it is between 0 and 1."""
    if len(img.shape) != 3:
        raise Exception("""Image must have 3 dimensions (channels x height x width). """
                        """Given {0}""".format(len(img.shape)))

    img_ch, _, _ = img.shape
    if img_ch != 3 and img_ch != 1:
        raise Exception("""Unsupported number of channels. """
                        """Must be 1 or 3, given {0}.""".format(img_ch))

    imgp = img
    if img_ch == 1:
        imgp = np.repeat(img, 3, axis=0)

    if not is_binary:
        imgp = imgp * 127.5 + 127.5
        imgp /= 255.

    return np.clip(imgp.transpose((1, 2, 0)), 0, 1)


def compose_imgs(a, b, is_a_binary=True, is_b_binary=False):
    """Place a and b side by side to be plotted."""
    ap = convert_to_rgb(a, is_binary=is_a_binary)
    bp = convert_to_rgb(b, is_binary=is_b_binary)

    if ap.shape != bp.shape:
        raise Exception("""A and B must have the same size. """
                        """{0} != {1}""".format(ap.shape, bp.shape))

    # ap.shape and bp.shape must have the same size here
    h, w, ch = ap.shape
    composed = np.zeros((h, 2*w, ch))
    composed[:, :w, :] = ap
    composed[:, w:, :] = bp

    return composed


def get_log_dir(log_dir, expt_name):
    """Compose the log_dir with the experiment name."""
    if log_dir is None:
        raise Exception('log_dir can not be None.')

    if expt_name is not None:
        return os.path.join(log_dir, expt_name)
    return log_dir


def mkdir(mypath):
    """Create a directory if it does not exist."""
    try:
        os.makedirs(mypath)
    except OSError as exc:
        if exc.errno == EEXIST and os.path.isdir(mypath):
            pass
        else:
            raise


def create_expt_dir(params):
    """Create the experiment directory and return it."""
    expt_dir = get_log_dir(params.log_dir, params.expt_name)

    # Create directories if they do not exist
    mkdir(params.log_dir)
    mkdir(expt_dir)

    # Save the parameters
    json.dump(params, open(os.path.join(expt_dir, 'params.json'), 'wb'),
              indent=4, sort_keys=True)

    return expt_dir


def plot_loss(loss, label, filename, log_dir):
    """Plot a loss function and save it in a file."""
    plt.figure(figsize=(5, 4))
    plt.plot(loss, label=label)
    plt.legend()
    plt.savefig(os.path.join(log_dir, filename))
    plt.clf()


def log(losses, atob, it_val, N=4, log_dir=DEFAULT_LOG_DIR, expt_name=None,
        is_a_binary=True, is_b_binary=False):
    """Log losses and atob results."""
    log_dir = get_log_dir(log_dir, expt_name)

    # Save the losses for further inspection
    pickle.dump(losses, open(os.path.join(log_dir, 'losses.pkl'), 'wb'))

    ###########################################################################
    #                             PLOT THE LOSSES                             #
    ###########################################################################
    plot_loss(losses['d'], 'discriminator', 'd_loss.png', log_dir)
    plot_loss(losses['d_val'], 'discriminator validation', 'd_val_loss.png', log_dir)

    plot_loss(losses['p2p'], 'Pix2Pix', 'p2p_loss.png', log_dir)
    plot_loss(losses['p2p_val'], 'Pix2Pix validation', 'p2p_val_loss.png', log_dir)

    ###########################################################################
    #                          PLOT THE A->B RESULTS                          #
    ###########################################################################
    plt.figure(figsize=(10, 6))
    for i in range(N*N):
        a, _ = next(it_val)

        bp = atob.predict(a)
        img = compose_imgs(a[0], bp[0], is_a_binary=is_a_binary, is_b_binary=is_b_binary)

        plt.subplot(N, N, i+1)
        plt.imshow(img)
        plt.axis('off')

    plt.savefig(os.path.join(log_dir, 'atob.png'))
    plt.clf()

    # Make sure all the figures are closed.
    plt.close('all')


def save_weights(models, log_dir=DEFAULT_LOG_DIR, expt_name=None):
    """Save the weights of the models into a file."""
    log_dir = get_log_dir(log_dir, expt_name)

    models.atob.save_weights(os.path.join(log_dir, ATOB_WEIGHTS_FILE), overwrite=True)
    models.d.save_weights(os.path.join(log_dir, D_WEIGHTS_FILE), overwrite=True)


def load_weights(atob, d, log_dir=DEFAULT_LOG_DIR, expt_name=None):
    """Load the weights into the corresponding models."""
    log_dir = get_log_dir(log_dir, expt_name)

    atob.load_weights(os.path.join(log_dir, ATOB_WEIGHTS_FILE))
    d.load_weights(os.path.join(log_dir, D_WEIGHTS_FILE))


def load_weights_of(m, weights_file, log_dir=DEFAULT_LOG_DIR, expt_name=None):
    """Load the weights of the model m."""
    log_dir = get_log_dir(log_dir, expt_name)

    m.load_weights(os.path.join(log_dir, weights_file))


def load_losses(log_dir=DEFAULT_LOG_DIR, expt_name=None):
    """Load the losses of the given experiment."""
    log_dir = get_log_dir(log_dir, expt_name)
    losses = pickle.load(open(os.path.join(log_dir, 'losses.pkl'), 'rb'))
    return losses


def load_params(params):
    """
    Load the parameters of an experiment and return them.

    The params passed as argument will be merged with the new params dict.
    If there is a conflict with a key, the params passed as argument prevails.
    """
    expt_dir = get_log_dir(params.log_dir, params.expt_name)

    expt_params = json.load(open(os.path.join(expt_dir, 'params.json'), 'rb'))

    # Update the loaded parameters with the current parameters. This will
    # override conflicting keys as expected.
    expt_params.update(params)

    return expt_params

[((11, 0, 11, 9), 'seaborn.set', 'sns.set', ({}, {}), '()', True, 'import seaborn as sns\n'), ((68, 15, 68, 37), 'numpy.zeros', 'np.zeros', ({(68, 24, 68, 36): '(h, 2 * w, ch)'}, {}), '((h, 2 * w, ch))', True, 'import numpy as np\n'), ((113, 4, 113, 30), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((114, 4, 114, 31), 'matplotlib.pyplot.plot', 'plt.plot', (), '', True, 'import matplotlib.pyplot as plt\n'), ((115, 4, 115, 16), 'matplotlib.pyplot.legend', 'plt.legend', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((117, 4, 117, 13), 'matplotlib.pyplot.clf', 'plt.clf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((140, 4, 140, 31), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((152, 4, 152, 13), 'matplotlib.pyplot.clf', 'plt.clf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((155, 4, 155, 20), 'matplotlib.pyplot.close', 'plt.close', ({(155, 14, 155, 19): '"""all"""'}, {}), "('all')", True, 'import matplotlib.pyplot as plt\n'), ((48, 15, 48, 40), 'numpy.repeat', 'np.repeat', (), '', True, 'import numpy as np\n'), ((81, 15, 81, 47), 'os.path.join', 'os.path.join', ({(81, 28, 81, 35): 'log_dir', (81, 37, 81, 46): 'expt_name'}, {}), '(log_dir, expt_name)', False, 'import os\n'), ((88, 8, 88, 27), 'os.makedirs', 'os.makedirs', ({(88, 20, 88, 26): 'mypath'}, {}), '(mypath)', False, 'import os\n'), ((116, 16, 116, 47), 'os.path.join', 'os.path.join', ({(116, 29, 116, 36): 'log_dir', (116, 38, 116, 46): 'filename'}, {}), '(log_dir, filename)', False, 'import os\n'), ((147, 8, 147, 30), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(147, 20, 147, 21): 'N', (147, 23, 147, 24): 'N', (147, 26, 147, 29): '(i + 1)'}, {}), '(N, N, i + 1)', True, 'import matplotlib.pyplot as plt\n'), ((148, 8, 148, 23), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(148, 19, 148, 22): 'img'}, {}), '(img)', True, 'import matplotlib.pyplot as plt\n'), ((149, 8, 149, 23), 'matplotlib.pyplot.axis', 'plt.axis', ({(149, 17, 149, 22): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((151, 16, 151, 49), 'os.path.join', 'os.path.join', ({(151, 29, 151, 36): 'log_dir', (151, 38, 151, 48): '"""atob.png"""'}, {}), "(log_dir, 'atob.png')", False, 'import os\n'), ((162, 29, 162, 69), 'os.path.join', 'os.path.join', ({(162, 42, 162, 49): 'log_dir', (162, 51, 162, 68): 'ATOB_WEIGHTS_FILE'}, {}), '(log_dir, ATOB_WEIGHTS_FILE)', False, 'import os\n'), ((163, 26, 163, 63), 'os.path.join', 'os.path.join', ({(163, 39, 163, 46): 'log_dir', (163, 48, 163, 62): 'D_WEIGHTS_FILE'}, {}), '(log_dir, D_WEIGHTS_FILE)', False, 'import os\n'), ((170, 22, 170, 62), 'os.path.join', 'os.path.join', ({(170, 35, 170, 42): 'log_dir', (170, 44, 170, 61): 'ATOB_WEIGHTS_FILE'}, {}), '(log_dir, ATOB_WEIGHTS_FILE)', False, 'import os\n'), ((171, 19, 171, 56), 'os.path.join', 'os.path.join', ({(171, 32, 171, 39): 'log_dir', (171, 41, 171, 55): 'D_WEIGHTS_FILE'}, {}), '(log_dir, D_WEIGHTS_FILE)', False, 'import os\n'), ((178, 19, 178, 54), 'os.path.join', 'os.path.join', ({(178, 32, 178, 39): 'log_dir', (178, 41, 178, 53): 'weights_file'}, {}), '(log_dir, weights_file)', False, 'import os\n'), ((105, 27, 105, 64), 'os.path.join', 'os.path.join', ({(105, 40, 105, 48): 'expt_dir', (105, 50, 105, 63): '"""params.json"""'}, {}), "(expt_dir, 'params.json')", False, 'import os\n'), ((126, 29, 126, 64), 'os.path.join', 'os.path.join', ({(126, 42, 126, 49): 'log_dir', (126, 51, 126, 63): '"""losses.pkl"""'}, {}), "(log_dir, 'losses.pkl')", False, 'import os\n'), ((184, 30, 184, 65), 'os.path.join', 'os.path.join', ({(184, 43, 184, 50): 'log_dir', (184, 52, 184, 64): '"""losses.pkl"""'}, {}), "(log_dir, 'losses.pkl')", False, 'import os\n'), ((197, 33, 197, 70), 'os.path.join', 'os.path.join', ({(197, 46, 197, 54): 'expt_dir', (197, 56, 197, 69): '"""params.json"""'}, {}), "(expt_dir, 'params.json')", False, 'import os\n'), ((90, 35, 90, 56), 'os.path.isdir', 'os.path.isdir', ({(90, 49, 90, 55): 'mypath'}, {}), '(mypath)', False, 'import os\n')]