#
#------------------------------------------------------------------------------
# Copyright (c) 2013-2014, Christian Therien
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#------------------------------------------------------------------------------
#
# test_eea.py - This file is part of the PySptools package.
#

"""
The following functions are tested:
    NFINDR
    ATGP
    PPI
    UCLS
    FCLS
    NNLS
"""

from __future__ import print_function

import os
import os.path as osp
import cProfile, pstats

import pysptools.util as util
import pysptools.eea as eea
import pysptools.abundance_maps as amp


_doProfile = False
def profile():
    if _doProfile == True:
        pr = cProfile.Profile()
        pr.enable()
        return pr


def stat(pr):
    if _doProfile == True:
        pr.disable()
        ps = pstats.Stats(pr)
        ps.strip_dirs()
        ps.sort_stats('time')
        ps.print_stats()


def test_PPI(data, wvl, result_path):
    print('Testing PPI')
    ppi = eea.PPI()
    # The format of the data is always (m x n x p),
    # for each class interface
    # the format of E is always (N x p)
    pr = profile()
    E = ppi.extract(data, 12, normalize=True)
    stat(pr)
    ppi.plot(result_path, wvl, suffix='test_PPI')
    test_amap(data, E, 'PPI', result_path, amaps='NNLS')


def test_NFINDR(data, wvl, result_path):
    print('Testing NFINDR')
    nfindr = eea.NFINDR()
    pr = profile()
    E = nfindr.extract(data, 12, maxit=5, normalize=True, ATGP_init=True)
    stat(pr)
    print('  Iterations:', nfindr.get_iterations())
    print('  End members indexes:', nfindr.get_idx())
    nfindr.plot(result_path, wvl, suffix='test_NFINDR')
    test_amap(data, E, 'NFINDR', result_path, amaps='NNLS')
    test_amap(data, E, 'NFINDR', result_path, amaps='FCLS')


def test_ATGP(data, wvl, result_path):
    print('Testing ATGP')
    atgp = eea.ATGP()
    pr = profile()
    E = atgp.extract(data, 12, normalize=True)
    stat(pr)
    atgp.plot(result_path, wvl, suffix='test_ATGP')
    #test_amap(data, E, 'ATGP', result_path, amaps='FCLS')
    test_amap(data, E, 'ATGP', result_path, amaps='NNLS')


def test_FIPPI(data, wvl, result_path):
    """
    Instantiate a Chein-I Chang kind of data processing flow
    HSI cube ->
    Virtual dimensionality ->
    PCA ->
    ATGP ->
    PPI ->
    NNLS
    """
    print('Testing FIPPI')
    fippi = eea.FIPPI()
    pr = profile()
    U = fippi.extract(data, 28, 1, normalize=True)
    fippi.plot(result_path, wvl, suffix='test_FIPPI')
    nnls = amp.NNLS()
    amap = nnls.map(data, U, normalize=True)
    stat(pr)
    nnls.plot(result_path, suffix='FIPPI')


def test_amap(data, U, umix_source, result_path, amaps=None):
    # if you normalize at the extract step, you need to
    # normalize at abundance map generation
    # and the opposite
    if amaps == None:
        test_UCLS(data, U, umix_source, result_path)
        test_NNLS(data, U, umix_source, result_path)
        test_FCLS(data, U, umix_source, result_path)
    else:
        if 'UCLS' in amaps: test_UCLS(data, U, umix_source, result_path)
        if 'NNLS' in amaps: test_NNLS(data, U, umix_source, result_path)
        if 'FCLS' in amaps: test_FCLS(data, U, umix_source, result_path)


def test_UCLS(data, U, umix_source, result_path):
    print('Testing UCLS')
    ucls = amp.UCLS()
    pr = profile()
    amap = ucls.map(data, U, normalize=True)
    stat(pr)
    ucls.plot(result_path, suffix=umix_source)


def test_NNLS(data, U, umix_source, result_path):
    print('Testing NNLS')
    nnls = amp.NNLS()
    pr = profile()
    amap = nnls.map(data, U, normalize=True)
    stat(pr)
    nnls.plot(result_path, colorMap='jet', suffix=umix_source)


def test_FCLS(data, U, umix_source, result_path):
    print('Testing FCLS')
    fcls = amp.FCLS()
    pr = profile()
    amap = fcls.map(data, U, normalize=True)
    stat(pr)
    fcls.plot(result_path, suffix=umix_source)


def tests_P2_7():
    data_path = '../data1'
    project_path = '../'
    result_path = osp.join(project_path, 'results')
    sample = 'SAMSON_mini.hdr'
    #sample = '92AV3C.hdr'
    #sample = 'hematite.hdr'

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

    data_file = osp.join(data_path, sample)
    data, wvl = util.load_ENVI_file(data_file)
    test_PPI(data, wvl, result_path)
    test_NFINDR(data, wvl, result_path)
    test_ATGP(data, wvl, result_path)
    test_FIPPI(data, wvl, result_path)


def tests_P3_3():
    import json
    import numpy as np
    data_path = r'..\data1'
    project_path = '..\\'
    result_path = osp.join(project_path, 'results')
    sample = 'SAMSON_mini'

    # 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:
        info = json.loads(content_file.read())

    test_PPI(data, info, result_path)
    test_NFINDR(data, info, result_path)
    test_ATGP(data, info, result_path)
    test_FIPPI(data, info, result_path)


def tests():
    import sys
    if sys.version_info[:2] == (2,7):
        tests_P2_7()
    if sys.version_info[:2] == (3,3):
        tests_P3_3()


if __name__ == '__main__':
    import sys
    print(sys.version_info)
    if sys.version_info[:2] == (2,7):
        tests_P2_7()
    if sys.version_info[:2] == (3,3):
        tests_P3_3()
