#!/usr/bin/env python3
"""Plot image produced by data generators in icenet/utils.py"""

import os
import argparse
import json
import random

import numpy as np
import matplotlib.pyplot as plt

from icenet import util
from plot import plot_sample
from itertools import islice


def plot_base_cnn_samples(outdir, samples, interactive=False):
    m = len(samples)
    minmax = util.get_minmax(samples)
    imgs, _ = list(islice(util.base_cnn_generator(
            samples, minmax, m, shuffle=False, verbose=True
    ), 1))[0]
    for i in range(m):
        print("Plotting sample %d/%d" % (i+1, m))
        sample = samples[i]
        img = imgs[i, :, :, :]

        fig, (ax1, ax2) = plt.subplots(1, 2, sharex='col', sharey='row')
        fig.suptitle('\n'.join([
            "Id: %s (rotated and cropped)" % sample.get('id'),
            "Iceberg? %s" % sample.get('is_iceberg', '-'),
            "Incident angle: %s" % repr(sample.get('inc_angle'))])
        )

        ax1.imshow(img[:, :, 0], vmin=0.2, vmax=0.8, aspect='auto')
        ax1.set_title("HH")
        ax2.imshow(img[:, :, 1], vmin=0.2, vmax=0.8, aspect='auto')
        ax2.set_title("HV")

        #fig.tight_layout()
        fig.subplots_adjust(top=0.80)
        plot_sample(outdir, sample, minmax, interactive=interactive)
        if not interactive:
            fig.savefig(os.path.join(outdir, "%s_base.png" % sample['id']))
        plt.close('all')


def plot_icenet_samples(outdir, samples, interactive=False):
    m = len(samples)
    minmax = util.get_minmax(samples)
    sections, y = list(islice(util.icenet_generator(
            samples, minmax, m, augment=False, verbose=True
    ), 1))[0]
    for i in range(m):
        print("Plotting sample %d/%d" % (i+1, m))
        sample = samples[i]

        for b in range(2):
            fig, axes = plt.subplots(3, 3, sharex='col', sharey='row')
            fig.suptitle('\n'.join([
                "Id: %s (rotated and cropped, band %d)" % (sample.get('id'), b+1),
                "Iceberg? %s" % sample.get('is_iceberg', '-'),
                "Incident angle: %s" % repr(sample.get('inc_angle'))])
            )

            for ridx in range(3):
                for cidx in range(3):
                    axes[ridx, cidx].imshow(
                            sections[ridx*3+cidx][i, :, :, b], 
                            vmin=0.2,
                            vmax=0.8,
                            aspect='auto'
                    )

            fig.subplots_adjust(top=0.80)
            if not interactive:
                fig.savefig(os.path.join(outdir, "%s_icenet_b%d.png" % (sample['id'], b)))

        plot_sample(outdir, sample, minmax, interactive=interactive)
        plt.close('all')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
            '-o',
            required=False,
            help="Output directory for samples."
    )
    parser.add_argument(
            '-g',
            required=True,
            choices=['base', 'icenet'],
            help="Generator to use (base|icenet)."
    )
    parser.add_argument(
            '-n',
            type=int,
            required=False,
            help="Number of samples to plot (randomized)"
    )
    parser.add_argument(
            '-i',
            required=False,
            action='store_true',
            help="Show plots (and don't store)"
    )
    parser.add_argument(
            'samples_file',
            help="JSON file with samples"
    )
    args = parser.parse_args()

    outdir = args.o or '.'
    if not os.path.isdir(outdir):
        raise Exception("Output directory does not exist")

    print("Loading samples ...")
    with open(args.samples_file) as f:
        samples = json.load(f)
    print("%d samples in set" % len(samples))

    if args.n:
        print("Pick %d random samples ..." % args.n)
        random.shuffle(samples)
        samples = samples[0:args.n]

    if args.g == 'base':
        plot_base_cnn_samples(outdir, samples, interactive=args.i or False)
    else:
        plot_icenet_samples(outdir, samples, interactive=args.i or False)

    print("Done.")