"""
Plot abundance maps stack for the methanol gas HSI cube.
"""

import os
import os.path as osp
import json
import numpy as np
import matplotlib.pyplot as plt
import pysptools.util as util
import pysptools.eea as eea
import pysptools.abundance_maps as amp


def get_endmembers(data, header, result_path):
    print('Endmembers extraction with NFINDR')
    nfindr = eea.NFINDR()
    U = nfindr.extract(data, 12, maxit=5, normalize=True, ATGP_init=True)
    nfindr.plot(result_path, header, suffix='gas')
    # return an array of endmembers
    return U


def gen_abundance_maps(data, U, result_path):
    print('Abundance maps generation with NNLS')
    nnls = amp.NNLS()
    amaps = nnls.map(data, U, normalize=True)
    nnls.plot(result_path, colorMap='jet', suffix='gas')
    # return an array of abundance maps
    return amaps


def plot_synthetic_gas_above(amap, colormap, result_path):
    print('Create and plot synthetic map for the gas_above')
    gas = (amap > 0.1) * amap
    stack = gas[:,:,4] + gas[:,:,5] + gas[:,:,9] + gas[:,:,10]
    plot_synthetic_image(stack, colormap, 'gas_above', result_path)


def plot_synthetic_gas_around(amap, colormap, result_path):
    print('Create and plot synthetic map for the gas_around')
    gas = (amap > 0.1) * amap
    stack = gas[:,:,6] + gas[:,:,7] + gas[:,:,8]
    plot_synthetic_image(stack, colormap, 'gas_around', result_path)


def plot_synthetic_burner(amap, colormap, result_path):
    print('Create and plot synthetic map for the burner')
    burner = (amap > 0.1) * amap
    stack = burner[:,:,2] + burner[:,:,3]
    plot_synthetic_image(stack, colormap, 'burner', result_path)


def plot_synthetic_image(image, colormap, desc, result_path):
    plt.ioff()
    img = plt.imshow(image)
    img.set_cmap(colormap)
    plt.colorbar()
    fout = osp.join(result_path, 'synthetic_{0}.png'.format(desc))
    plt.savefig(fout)
    plt.clf()


if __name__ == '__main__':
    data_path = '../data1'
    project_path = '../'
    result_path = osp.join(project_path, 'results')
    sample = 'burner'

    fin = osp.join(data_path, sample)
    if osp.exists(result_path) == False:
        os.makedirs(result_path)

    # load the cube
    data_file = osp.join(data_path, sample+'.jdata')
    with open(data_file, 'r') as content_file:
        data = np.array(json.loads(content_file.read()))
    info_file = osp.join(data_path, sample+'.jhead')
    with open(info_file, 'r') as content_file:
        header = json.loads(content_file.read())

    # Telops cubes are flipped left-right
    # Flipping them again restore the orientation
    data = np.fliplr(data)

    U = get_endmembers(data, header, result_path)
    amaps = gen_abundance_maps(data, U, result_path)
    # best color maps: hot and spectral
    plot_synthetic_gas_above(amaps, 'spectral', result_path)
    plot_synthetic_gas_around(amaps, 'hot', result_path)
    plot_synthetic_burner(amaps, 'hot', result_path)
