"""Loader for the Labeled Faces in the Wild (LFW) dataset

This dataset is a collection of JPEG pictures of famous people collected
over the internet, all details are available on the official website:

    http://vis-www.cs.umass.edu/lfw/

Each picture is centered on a single face. The typical task is called
Face Verification: given a pair of two pictures, a binary classifier
must predict whether the two images are from the same person.

An alternative task, Face Recognition or Face Identification is:
given the picture of the face of an unknown person, identify the name
of the person by referring to a gallery of previously seen pictures of
identified persons.

Both Face Verification and Face Recognition are tasks that are typically
performed on the output of a model trained to perform Face Detection. The
most popular model for Face Detection is called Viola-Johns and is
implemented in the OpenCV library. The LFW faces were extracted by this face
detector from various online websites.
"""
# Copyright (c) 2011 Olivier Grisel <olivier.grisel@ensta.org>
# License: BSD 3 clause

from os import listdir, makedirs, remove
from os.path import join, exists, isdir

from sklearn.utils import deprecated

import logging
import numpy as np

try:
    import urllib.request as urllib  # for backwards compatibility
except ImportError:
    import urllib

from .base import get_data_home, Bunch
from ..externals.joblib import Memory

from ..externals.six import b

logger = logging.getLogger(__name__)

BASE_URL = "http://vis-www.cs.umass.edu/lfw/"
ARCHIVE_NAME = "lfw.tgz"
FUNNELED_ARCHIVE_NAME = "lfw-funneled.tgz"
TARGET_FILENAMES = [
    'pairsDevTrain.txt',
    'pairsDevTest.txt',
    'pairs.txt',
]

def scale_face(face):
    """Scale back to 0-1 range in case of normalization for plotting"""
    scaled = face - face.min()
    scaled /= scaled.max()
    return scaled

#
# Common private utilities for data fetching from the original LFW website
# local disk caching, and image decoding.
#

def check_fetch_lfw(data_home=None, funneled=True, download_if_missing=True):
    """Helper function to download any missing LFW data"""
    data_home = get_data_home(data_home=data_home)
    lfw_home = join(data_home, "lfw_home")

    if funneled:
        archive_path = join(lfw_home, FUNNELED_ARCHIVE_NAME)
        data_folder_path = join(lfw_home, "lfw_funneled")
        archive_url = BASE_URL + FUNNELED_ARCHIVE_NAME
    else:
        archive_path = join(lfw_home, ARCHIVE_NAME)
        data_folder_path = join(lfw_home, "lfw")
        archive_url = BASE_URL + ARCHIVE_NAME

    if not exists(lfw_home):
        makedirs(lfw_home)

    for target_filename in TARGET_FILENAMES:
        target_filepath = join(lfw_home, target_filename)
        if not exists(target_filepath):
            if download_if_missing:
                url = BASE_URL + target_filename
                logger.warning("Downloading LFW metadata: %s", url)
                urllib.urlretrieve(url, target_filepath)
            else:
                raise IOError("%s is missing" % target_filepath)

    if not exists(data_folder_path):

        if not exists(archive_path):
            if download_if_missing:
                logger.warning("Downloading LFW data (~200MB): %s", archive_url)
                urllib.urlretrieve(archive_url, archive_path)
            else:
                raise IOError("%s is missing" % target_filepath)

        import tarfile
        logger.info("Decompressing the data archive to %s", data_folder_path)
        tarfile.open(archive_path, "r:gz").extractall(path=lfw_home)
        remove(archive_path)

    return lfw_home, data_folder_path

def _load_imgs(file_paths, slice_, color, resize):
    """Internally used to load images"""

    # Try to import imread and imresize from PIL. We do this here to prevent
    # the whole sklearn.datasets module from depending on PIL.
    try:
        try:
            from scipy.misc import imread
        except ImportError:
            from scipy.misc.pilutil import imread
        from scipy.misc import imresize
    except ImportError:
        raise ImportError("The Python Imaging Library (PIL)"
                          " is required to load data from jpeg files")

    # compute the portion of the images to load to respect the slice_ parameter
    # given by the caller
    default_slice = (slice(0, 250), slice(0, 250))
    if slice_ is None:
        slice_ = default_slice
    else:
        slice_ = tuple(s or ds for s, ds in zip(slice_, default_slice))

    h_slice, w_slice = slice_
    h = (h_slice.stop - h_slice.start) // (h_slice.step or 1)
    w = (w_slice.stop - w_slice.start) // (w_slice.step or 1)

    if resize is not None:
        resize = float(resize)
        h = int(resize * h)
        w = int(resize * w)

    # allocate some contiguous memory to host the decoded image slices
    n_faces = len(file_paths)
    if not color:
        faces = np.zeros((n_faces, h, w), dtype=np.float32)
    else:
        faces = np.zeros((n_faces, h, w, 3), dtype=np.float32)

    # iterate over the collected file path to load the jpeg files as numpy
    # arrays
    for i, file_path in enumerate(file_paths):
        if i % 1000 == 0:
            logger.info("Loading face #%05d / %05d", i + 1, n_faces)

        # Checks if jpeg reading worked. Refer to issue #3594 for more
        # details.
        img = imread(file_path)
        if img.ndim is 0:
            raise RuntimeError("Failed to read the image file %s, "
                               "Please make sure that libjpeg is installed"
                               % file_path)

        face = np.asarray(img[slice_], dtype=np.float32)
        face /= 255.0  # scale uint8 coded colors to the [0.0, 1.0] floats
        if resize is not None:
            face = imresize(face, resize)
        if not color:
            # average the color channels to compute a gray levels
            # representaion
            face = face.mean(axis=2)

        faces[i, ...] = face

    return faces

#
# Task #1:  Face Identification on picture with names
#

def _fetch_lfw_people(data_folder_path, slice_=None, color=False, resize=None,
                      min_faces_per_person=0):
    """Perform the actual data loading for the lfw people dataset

    This operation is meant to be cached by a joblib wrapper.
    """
    # scan the data folder content to retain people with more that
    # `min_faces_per_person` face pictures
    person_names, file_paths = [], []
    for person_name in sorted(listdir(data_folder_path)):
        folder_path = join(data_folder_path, person_name)
        if not isdir(folder_path):
            continue
        paths = [join(folder_path, f) for f in listdir(folder_path)]
        n_pictures = len(paths)
        if n_pictures >= min_faces_per_person:
            person_name = person_name.replace('_', ' ')
            person_names.extend([person_name] * n_pictures)
            file_paths.extend(paths)

    n_faces = len(file_paths)
    if n_faces == 0:
        raise ValueError("min_faces_per_person=%d is too restrictive" %
                         min_faces_per_person)

    target_names = np.unique(person_names)
    target = np.searchsorted(target_names, person_names)

    faces = _load_imgs(file_paths, slice_, color, resize)

    # shuffle the faces with a deterministic RNG scheme to avoid having
    # all faces of the same person in a row, as it would break some
    # cross validation and learning algorithms such as SGD and online
    # k-means that make an IID assumption

    indices = np.arange(n_faces)
    np.random.RandomState(42).shuffle(indices)
    faces, target = faces[indices], target[indices]
    return faces, target, target_names

def fetch_lfw_people(data_home=None, funneled=True, resize=0.5,
                     min_faces_per_person=0, color=False,
                     slice_=(slice(70, 195), slice(78, 172)),
                     download_if_missing=True):
    """Loader for the Labeled Faces in the Wild (LFW) people dataset

    This dataset is a collection of JPEG pictures of famous people
    collected on the internet, all details are available on the
    official website:

        http://vis-www.cs.umass.edu/lfw/

    Each picture is centered on a single face. Each pixel of each channel
    (color in RGB) is encoded by a float in range 0.0 - 1.0.

    The task is called Face Recognition (or Identification): given the
    picture of a face, find the name of the person given a training set
    (gallery).

    The original images are 250 x 250 pixels, but the default slice and resize
    arguments reduce them to 62 x 74.

    Parameters
    ----------
    data_home : optional, default: None
        Specify another download and cache folder for the datasets. By default
        all scikit learn data is stored in '~/scikit_learn_data' subfolders.

    funneled : boolean, optional, default: True
        Download and use the funneled variant of the dataset.

    resize : float, optional, default 0.5
        Ratio used to resize the each face picture.

    min_faces_per_person : int, optional, default None
        The extracted dataset will only retain pictures of people that have at
        least `min_faces_per_person` different pictures.

    color : boolean, optional, default False
        Keep the 3 RGB channels instead of averaging them to a single
        gray level channel. If color is True the shape of the data has
        one more dimension than than the shape with color = False.

    slice_ : optional
        Provide a custom 2D slice (height, width) to extract the
        'interesting' part of the jpeg files and avoid use statistical
        correlation from the background

    download_if_missing : optional, True by default
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.

    Returns
    -------
    dataset : dict-like object with the following attributes:

    dataset.data : numpy array of shape (13233, 2914)
        Each row corresponds to a ravelled face image of original size 62 x 47
        pixels. Changing the ``slice_`` or resize parameters will change the shape
        of the output.

    dataset.images : numpy array of shape (13233, 62, 47)
        Each row is a face image corresponding to one of the 5749 people in
        the dataset. Changing the ``slice_`` or resize parameters will change the shape
        of the output.

    dataset.target : numpy array of shape (13233,)
        Labels associated to each face image. Those labels range from 0-5748
        and correspond to the person IDs.

    dataset.DESCR : string
        Description of the Labeled Faces in the Wild (LFW) dataset.
    """
    lfw_home, data_folder_path = check_fetch_lfw(
        data_home=data_home, funneled=funneled,
        download_if_missing=download_if_missing)
    logger.info('Loading LFW people faces from %s', lfw_home)

    # wrap the loader in a memoizing function that will return memmaped data
    # arrays for optimal memory usage
    m = Memory(cachedir=lfw_home, compress=6, verbose=0)
    load_func = m.cache(_fetch_lfw_people)

    # load and memoize the pairs as np arrays
    faces, target, target_names = load_func(
        data_folder_path, resize=resize,
        min_faces_per_person=min_faces_per_person, color=color, slice_=slice_)

    # pack the results as a Bunch instance
    return Bunch(data=faces.reshape(len(faces), -1), images=faces,
                 target=target, target_names=target_names,
                 DESCR="LFW faces dataset")

#
# Task #2:  Face Verification on pairs of face pictures
#

def _fetch_lfw_pairs(index_file_path, data_folder_path, slice_=None,
                     color=False, resize=None):
    """Perform the actual data loading for the LFW pairs dataset

    This operation is meant to be cached by a joblib wrapper.
    """
    # parse the index file to find the number of pairs to be able to allocate
    # the right amount of memory before starting to decode the jpeg files
    with open(index_file_path, 'rb') as index_file:
        split_lines = [ln.strip().split(b('\t')) for ln in index_file]
    pair_specs = [sl for sl in split_lines if len(sl) > 2]
    n_pairs = len(pair_specs)

    # interating over the metadata lines for each pair to find the filename to
    # decode and load in memory
    target = np.zeros(n_pairs, dtype=np.int)
    file_paths = list()
    for i, components in enumerate(pair_specs):
        if len(components) == 3:
            target[i] = 1
            pair = (
                (components[0], int(components[1]) - 1),
                (components[0], int(components[2]) - 1),
            )
        elif len(components) == 4:
            target[i] = 0
            pair = (
                (components[0], int(components[1]) - 1),
                (components[2], int(components[3]) - 1),
            )
        else:
            raise ValueError("invalid line %d: %r" % (i + 1, components))
        for j, (name, idx) in enumerate(pair):
            try:
                person_folder = join(data_folder_path, name)
            except TypeError:
                person_folder = join(data_folder_path, str(name, 'UTF-8'))
            filenames = list(sorted(listdir(person_folder)))
            file_path = join(person_folder, filenames[idx])
            file_paths.append(file_path)

    pairs = _load_imgs(file_paths, slice_, color, resize)
    shape = list(pairs.shape)
    n_faces = shape.pop(0)
    shape.insert(0, 2)
    shape.insert(0, n_faces // 2)
    pairs.shape = shape

    return pairs, target, np.array(['Different persons', 'Same person'])

@deprecated("Function 'load_lfw_people' has been deprecated in 0.17 and will be "
            "removed in 0.19."
            "Use fetch_lfw_people(download_if_missing=False) instead.")
def load_lfw_people(download_if_missing=False, **kwargs):
    """Alias for fetch_lfw_people(download_if_missing=False)

    Check fetch_lfw_people.__doc__ for the documentation and parameter list.
    """
    return fetch_lfw_people(download_if_missing=download_if_missing, **kwargs)

def fetch_lfw_pairs(subset='train', data_home=None, funneled=True, resize=0.5,
                    color=False, slice_=(slice(70, 195), slice(78, 172)),
                    download_if_missing=True):
    """Loader for the Labeled Faces in the Wild (LFW) pairs dataset

    This dataset is a collection of JPEG pictures of famous people
    collected on the internet, all details are available on the
    official website:

        http://vis-www.cs.umass.edu/lfw/

    Each picture is centered on a single face. Each pixel of each channel
    (color in RGB) is encoded by a float in range 0.0 - 1.0.

    The task is called Face Verification: given a pair of two pictures,
    a binary classifier must predict whether the two images are from
    the same person.

    In the official `README.txt`_ this task is described as the
    "Restricted" task.  As I am not sure as to implement the
    "Unrestricted" variant correctly, I left it as unsupported for now.

      .. _`README.txt`: http://vis-www.cs.umass.edu/lfw/README.txt

    The original images are 250 x 250 pixels, but the default slice and resize
    arguments reduce them to 62 x 74.

    Read more in the :ref:`User Guide <labeled_faces_in_the_wild>`.

    Parameters
    ----------
    subset : optional, default: 'train'
        Select the dataset to load: 'train' for the development training
        set, 'test' for the development test set, and '10_folds' for the
        official evaluation set that is meant to be used with a 10-folds
        cross validation.

    data_home : optional, default: None
        Specify another download and cache folder for the datasets. By
        default all scikit learn data is stored in '~/scikit_learn_data'
        subfolders.

    funneled : boolean, optional, default: True
        Download and use the funneled variant of the dataset.

    resize : float, optional, default 0.5
        Ratio used to resize the each face picture.

    color : boolean, optional, default False
        Keep the 3 RGB channels instead of averaging them to a single
        gray level channel. If color is True the shape of the data has
        one more dimension than than the shape with color = False.

    slice_ : optional
        Provide a custom 2D slice (height, width) to extract the
        'interesting' part of the jpeg files and avoid use statistical
        correlation from the background

    download_if_missing : optional, True by default
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.

    Returns
    -------
    The data is returned as a Bunch object with the following attributes:

    data : numpy array of shape (2200, 5828)
        Each row corresponds to 2 ravel'd face images of original size 62 x 47
        pixels. Changing the ``slice_`` or resize parameters will change the shape
        of the output.

    pairs : numpy array of shape (2200, 2, 62, 47)
        Each row has 2 face images corresponding to same or different person
        from the dataset containing 5749 people. Changing the ``slice_`` or resize
        parameters will change the shape of the output.

    target : numpy array of shape (13233,)
        Labels associated to each pair of images. The two label values being
        different persons or the same person.

    DESCR : string
        Description of the Labeled Faces in the Wild (LFW) dataset.

    """
    lfw_home, data_folder_path = check_fetch_lfw(
        data_home=data_home, funneled=funneled,
        download_if_missing=download_if_missing)
    logger.info('Loading %s LFW pairs from %s', subset, lfw_home)

    # wrap the loader in a memoizing function that will return memmaped data
    # arrays for optimal memory usage
    m = Memory(cachedir=lfw_home, compress=6, verbose=0)
    load_func = m.cache(_fetch_lfw_pairs)

    # select the right metadata file according to the requested subset
    label_filenames = {
        'train': 'pairsDevTrain.txt',
        'test': 'pairsDevTest.txt',
        '10_folds': 'pairs.txt',
    }
    if subset not in label_filenames:
        raise ValueError("subset='%s' is invalid: should be one of %r" % (
            subset, list(sorted(label_filenames.keys()))))
    index_file_path = join(lfw_home, label_filenames[subset])

    # load and memoize the pairs as np arrays
    pairs, target, target_names = load_func(
        index_file_path, data_folder_path, resize=resize, color=color,
        slice_=slice_)

    # pack the results as a Bunch instance
    return Bunch(data=pairs.reshape(len(pairs), -1), pairs=pairs,
                 target=target, target_names=target_names,
                 DESCR="'%s' segment of the LFW pairs dataset" % subset)

@deprecated("Function 'load_lfw_pairs' has been deprecated in 0.17 and will be "
            "removed in 0.19."
            "Use fetch_lfw_pairs(download_if_missing=False) instead.")
def load_lfw_pairs(download_if_missing=False, **kwargs):
    """Alias for fetch_lfw_pairs(download_if_missing=False)

    Check fetch_lfw_pairs.__doc__ for the documentation and parameter list.
    """
    return fetch_lfw_pairs(download_if_missing=download_if_missing, **kwargs)


from __future__ import print_function
import argparse
import importlib
import os
from shutil import rmtree
from six.moves import input
import sys
import textwrap
import time
import json

# TODO:
#       catch and log exceptions in examples files that fail to open

DIRECTORIES = {
    'file'    : '../../examples/plotting/file',
    'notebook': '../../examples/plotting/notebook',
    'server'  : '../../examples/plotting/server',
    'ggplot'  : '../../examples/compat/ggplot',
    'glyphs'  : '../../examples/glyphs',
    'mpl'     : '../../examples/compat/mpl',
    'pandas'  : '../../examples/compat/pandas',
    'seaborn' : '../../examples/compat/seaborn',
    'charts'  : '../../examples/charts',
}

DEFAULT_TEST_FILES = [
    '../../examples/plotting/file/stocks.py',
    '../../examples/plotting/file/glucose.py',
    '../../examples/compat/ggplot/density.py',
    '../../examples/plotting/server/stocks.py',
    '../../examples/plotting/server/glucose.py',
    '../../examples/plotting/notebook/candlestick.ipynb',
    '../../examples/plotting/notebook/glucose.ipynb',
    '../../examples/compat/seaborn/violin.py',
    '../../examples/charts/boxplot.py',
]

SESSION_FILE = os.path.abspath("INTERACTIVE_TESTER_SESSION.json")

def get_parser():
    """Create the parser that will be used to add arguments to the script.
    """

    parser = argparse.ArgumentParser(description=textwrap.dedent("""
                    Tests a selection of .py or .ipynb bokeh example files.

                    The --location option allows you to select a specific examples subdirectory to test all files in,
                    ignoring __init__.py

                    Location arguments can be any valid path to a folder with the examples, like:

                     -l /path/to/my/examplesyou can choose:

                    or any of the pre-built keywords that point to the related examples:
                        - file
                        - notebook
                        - server
                        - ggplot
                        - glyphs
                        - mpl
                        - pandas
                        - seaborn
                    """), formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('--no-log', action='store_true', dest='nolog', default=False,
                        help="don't save a log of any errors discovered")
    parser.add_argument('-l', '--location', action='store', default=False,
                        help="example directory in which you wish to test")
    parser.add_argument('--reuse-session', action='store_true', dest='reuseSession', default=False,
                        help="do not clean last session log and start from where you left")
    parser.add_argument('--notebook-options', action='store', dest='notebookOptions', default="",
                        help="options to be forwarded to ipython notebook to customize it's behaviour")

    return parser

def depend_check(dependency):
    """
    Make sure a given dependency is installed
    """

    try:
        importlib.import_module(dependency)
        found = True
    except ImportError as e:
        print("%s\nPlease use conda or pip to install the necessary dependency." % (e))
        found = False

    return found

def save_session(session):
    """
    Save the session object to the SESSION_FILE

    Args:
        session(dict): dict with all the example files and results of each run
    """
    with open(SESSION_FILE, 'w') as res_file:
        json.dump(session, res_file)

def get_session():
    """
    Return last stored session
    """
    try:
        with open(SESSION_FILE, 'r') as res_file:
            return json.load(res_file)
    except IOError:
        return {}

def clean_session():
    """
    Removes previous session file
    """
    if os.path.exists(SESSION_FILE):
        os.remove(SESSION_FILE)

def main(testing_ground=None, notebook_options=""):
    """
    Collect and run .py or .ipynb examples from a set list or given examples directory, ignoring __init__.py
    User input is collected to determine a properly or improperly displayed page

    """

    # Create a testing directory if one does not exist, then cd into it

    testing_directory = 'tmp_test'

    if not os.path.exists(testing_directory):
        os.mkdir(testing_directory)

    os.chdir(testing_directory)

    if testing_ground:
        log_name = results.location

        TestFiles = [
            fileName for fileName in os.listdir('%s/.' % testing_ground)
            if fileName.endswith(('.py', '.ipynb')) and fileName != '__init__.py'
        ]

    else:
        log_name = "fast"

        TestFiles = DEFAULT_TEST_FILES

    Log = []

    lastSession = get_session()

    for index, fileName in enumerate(TestFiles):
        if testing_ground:
            fileName = "%s/%s" % (testing_ground, fileName)
        try:

            if not fileName in lastSession:
                lastSession[fileName] = "TESTING..."
                save_session(lastSession)

                command = get_cmd(fileName, notebook_options)
                opener(fileName, command)

                if results.nolog:
                    # Don't display 'next file' message after opening final file in a dir
                    if index != len(TestFiles)-1:
                        input("\nPress enter to open next file ")
                else:
                    ErrorReport = test_status()
                    if ErrorReport:
                        Log.append("\n\n%s: \n %s" % (fileName, ErrorReport))

                    lastSession[fileName] = ErrorReport
                    save_session(lastSession)
            else:
                prevRes = lastSession[fileName]
                if prevRes == "TESTING...":
                    print("RESULT OF %s LAST RUN NOT REGISTERED!!" % fileName)
                    ErrorReport = test_status()
                    lastSession[fileName] = ErrorReport
                    save_session(lastSession)
                else:
                    print("%s detected in last session: SKIPPING" % fileName)

        except (KeyboardInterrupt, EOFError):
            break

    # exit the testing directory and delete it

    os.chdir('../')
    rmtree(testing_directory)

    if Log:
        logger(Log, log_name)

def get_cmd(some_file, notebook_options=""):
    """Determines how to open a file depending
    on whether it is a .py or a .ipynb file
    """

    if some_file.endswith('.py'):
        command = "python"
    elif some_file.endswith('.ipynb'):
        command = "ipython notebook %s" % notebook_options

    return command

def opener(some_file, command):
    """Print to screen what file is being opened and then open the file using
    the command method provided.
    """

    print("\nOpening %s\n" % some_file.strip('../'))
    os.system("%s %s" % (command, some_file))

def test_status():
    """Collect user input to determine if a file displayed correctly or incorrectly.
    In the case of incorrectly displayed plots, an 'ErrorReport' string is returned.
    """

    status = input("Did the plot(s) display correctly? (y/n) ")
    while not status.startswith(('y', 'n')):
        print("")
        status = input("Unexpected answer. Please type y or n. ")
    if status.startswith('n'):
        ErrorReport = input("Please describe the problem: ")
        return ErrorReport

def logger(error_array, name):
    """
    Log errors by appending to a .txt file.  The name and directory the file is saved into
    is provided by the name and log_dir args.
    """

    logfile = "%s_examples_testlog.txt" % name
    if os.path.exists(logfile):
        os.remove(logfile)

    with open(logfile, 'a') as f:
        print("")
        print("\nWriting error log to %s" % logfile)
        for error in error_array:
            f.write("%s\n" % error)

if __name__ == '__main__':

    if not depend_check('bokeh'):
        sys.exit(1)

    parser = get_parser()
    results = parser.parse_args()

    if results.location:
        if results.location and results.location in DIRECTORIES:
            target = results.location

            if target in ['ggplot', 'pandas', 'seaborn', 'charts']:
                if not depend_check(target):
                    sys.exit(1)

            test_dir = DIRECTORIES[target]

        elif os.path.exists(results.location):
            # in case target is not one of the recognized keys and is a
            # valid path we can run the examples in that folder
            test_dir = results.location
            print("Running examples in custom location:", test_dir)
        else:
            print("Test location '%s' not recognized.\nPlease type 'python interactive_tester.py -h' for a list of valid test directories."
                 % results.location)
            sys.exit(1)
    else:
        test_dir = None

    if results.location == 'server' or test_dir is None:
        print("Server examples require bokeh-server. Make sure you've typed 'bokeh-server' in another terminal tab.")
        time.sleep(5)

    if not results.reuseSession:
        print("cleaning previous session file...",)
        clean_session()
        print("OK")

    main(test_dir, notebook_options=results.notebookOptions)


import matplotlib.pyplot as plt
from numpy import *

###List of variables
#	r_in		[10**10 cm]			innder radius
#	r_out		[10**10 cm]			outer radius
#	step		[10**10 cm]			step of plot
#	alfa		[]					parameter of accretion
#	M_16		[10**16 g.s**(-1)]	accretion flow
#	m_1			[solar mass]		mass of compact object
#	R_hv		[10**10 cm]			radius of compact object
#	R_10		[10**10 cm]			distance from compact object
#	f								numerical factor

###List of computed parameters
#	Surface density					[g.cm**(-2)]		(sigma)
#	Height							[cm]				(H)
#	Density							[g.cm**(-3)]		(rho)
#	Central disc temeprature		[K]					(T_c)
#	Opacity							[]					(tau)
#	viscosity						[cm**2.s**(-1)]		(nu)
#	radial velocity towards center	[cm.s**(-1)]		(v_r)

###function solutions parameters
# 	parameter 1		r_in
# 	parameter 2		r_out
# 	parameter 3		step
# 	parameter 4		alfa
# 	parameter 5		M_16
# 	parameter 6		m_1
# 	parameter 7		R_hv

def solutions(r_in,r_out,step,alfa,M_16,m_1,R_hv):
	#defining lists
	list_function = arange(r_in,r_out,step)
	R_10_l,surface_density_l,height_l,density_l,Fx = ([] for i in range(5))
	temperature_l,opacity_l,viscosity_l,radial_velocity_l = ([] for i in range(4))
	#computation and appending to lists
	for R_10 in list_function:
		f=(1-((R_hv)/(R_10))**(1.0/2))**(1.0/4)
		surface_density = 5.2*alfa**(-4.0/5)*M_16**(7.0/10)*m_1**(1.0/4)*R_10**(-3.0/4)*f**(14.0/5)
		height =  1.7*10**8*alfa**(-1.0/10)*M_16**(3.0/20)*m_1**(-3.0/8)*R_10**(9.0/8)*f**(3.0/5)
		density = 3.1*10**(-8)*alfa**(-7.0/10)*M_16**(11.0/20)*m_1**(5.0/8)*R_10**(-15.0/8)*f**(11.0/5)
		temperature = 1.4*10**4*alfa**(-1.0/5)*M_16**(3.0/10)*m_1**(1.0/4)*R_10**(-3.0/4)*f**(6.0/5)
		opacity = 190*alfa**(-4.0/5)*M_16**(1.0/5)*f**(4.0/5)
		viscosity = 1.8*10**14*alfa**(4.0/5)*M_16**(3.0/10)*m_1**(-1.0/4)*R_10**(3.0/4)*f**(6.0/5)
		radial_velocity = 2.7*10**4*alfa**(4.0/5)*M_16**(3.0/10)*m_1**(-1.0/4)*R_10**(-1.0/4)*f**(-14.0/5)
		R_10_l.append(R_10)
		surface_density_l.append(surface_density)
		height_l.append(height)
		density_l.append(density)
		temperature_l.append(temperature)
		opacity_l.append(opacity)
		viscosity_l.append(viscosity)
		radial_velocity_l.append(radial_velocity)
		Fx.append(f)
	#transformation R_10 to kolimeters
	R_km = [ x / 10**(-4) for x in R_10_l]
	return R_km, surface_density_l, height_l, density_l,temperature_l,opacity_l,viscosity_l,radial_velocity_l,Fx

#for definitions of parameters look up
r_in =1.0001*10**(-4)
r_out =10**(-2)
step = 10**(-6)
alfa = 0.5
M_16 = 63
m_1 = 1.5
R_hv = 1.0*10**(-4)

lists=solutions(r_in,r_out,step,alfa,M_16,m_1,R_hv)

print 30*"-"
print "Used parameter values"
print 30*"-"
print "innder radius:", 10*".",r_in, 10*".", "[10$^{10}$ cm]"
print "outer radius:", 10*".", r_out, 10*".", "[10$^{10}$ cm]"
print "step of plot:", 10*".", step, 10*".", "[10$^{10}$ cm]"
print "parameter of accretion alfa:", 10*".", alfa
print "accretion flow:", 10*".", M_16, 10*".", "[10$^6$ g.s${-1)}$]"
print "mass of compact object:", 10*".", m_1, 10*".", "[solar mass]"
print "radius of compact object:", 10*".", R_hv, 10*".", "[10$^{10}$ cm]"

plt.plot(lists[0], lists[1])
plt.title('surface density')
plt.xlabel('radius [km]')
plt.ylabel('surface density [g.cm$^{-2}$] ')
plt.grid()
plt.savefig("surface density")
plt.gcf().clear()

plt.plot(lists[0], lists[2])
plt.title('height')
plt.xlabel('radius [km]')
plt.ylabel('height [cm] ')
plt.grid()
plt.savefig("height")
plt.gcf().clear()

plt.plot(lists[0], lists[3])
plt.title('density')
plt.xlabel('radius [km]')
plt.ylabel('density [g.cm$^{-3}$] ')
plt.grid()
plt.savefig("density")
plt.gcf().clear()

plt.plot(lists[0], lists[4])
plt.title('temperature')
plt.xlabel('radius [km]')
plt.ylabel('temperature [K] ')
plt.grid()
plt.savefig("temperature")
plt.gcf().clear()

plt.plot(lists[0], lists[5])
plt.title('opacity')
plt.xlabel('radius [km]')
plt.ylabel('opacity ')
plt.grid()
plt.savefig("opacity")
plt.gcf().clear()

plt.plot(lists[0], lists[6])
plt.title('viscosity')
plt.xlabel('radius [km]')
plt.ylabel('viscosity [cm$^{2}$.s$^{-1}$] ')
plt.grid()
plt.savefig("viscosity")
plt.gcf().clear()

plt.plot(lists[0], lists[7])
plt.title('radial velocity')
plt.xlabel('radius [km]')
plt.ylabel('radial velocity [cm.s$^{-1}$] ')
plt.grid()
plt.savefig("radial velocity")
plt.gcf().clear()


# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
import numpy as np
import sys
import astropy.io.fits as afits
import collections
import tdose_utilities as tu
import tdose_extract_spectra as tes
import tdose_build_mock_cube as tbmc
import pdb
import scipy.ndimage.filters as snf
import matplotlib as mpl
mpl.use('Agg') # prevent pyplot from opening window; enables closing ssh session with detached screen running TDOSE
import matplotlib.pyplot as plt
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectra(model_cube_file,source_association_dictionary=None,nameext='tdose_spectrum',outputdir='./',clobber=False,
                    variance_cube_file=None,variance_cube_ext='ERROR',source_model_cube_file=None,source_cube_ext='DATA',
                    model_cube_ext='DATA',layer_scale_ext='WAVESCL',data_cube_file=None,verbose=True):
    """
    Assemble the spectra determined by the wavelength layer scaling of the normalized models
    when generating the source model cube

    --- INPUT ---
    model_cube_file                     Model cube to base extraction on (using header info and layer scales)
    source_association_dictionary       Source association dictionary defining what sources should be combined into
                                        objects (individual spectra).
    nameext                             The name extension to use for saved spectra
    outputdir                           Directory to save spectra to
    clobber                             Overwrite spectra if they already exists
    variance_cube_file                  File containing variance cube of data to be used to estimate nois on 1D spectrum
    variance_cube_ext                   Extension of variance cube to use
    source_model_cube_file              The source model cube defining the individual sources
    source_cube_ext                     Extension of source model cube file that contins source models
    model_cube_ext                      Extension of model cube file that contains model
    layer_scale_ext                     Extension of model cube file that contains the layer scales
    data_cube_file                      File containing original data cube used for extraction of aperture spectra
    verbose

    --- EXAMPLE OF USE ---

    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Loading data needed for spectral assembly')
    model_cube        = afits.open(model_cube_file)[model_cube_ext].data
    model_cube_hdr    = afits.open(model_cube_file)[model_cube_ext].header
    layer_scale_arr   = afits.open(model_cube_file)[layer_scale_ext].data
    if variance_cube_file is not None:
        stddev_cube       = np.sqrt(afits.open(variance_cube_file)[variance_cube_ext].data) # turn varinace into standard deviation
        source_model_cube = afits.open(source_model_cube_file)[source_cube_ext].data
    else:
        stddev_cube       = None
        source_model_cube = None
    Nsources = layer_scale_arr.shape[0]

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if data_cube_file is not None:
        if verbose: print(' - Loading data cube ')
        data_cube  = afits.open(data_cube_file)[model_cube_ext].data
    else:
        data_cube  = None

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if source_association_dictionary is None:
        if verbose: print(' - Building default source association dictionary ' \
                          '(determining what sources are combined into objects), i.e., one source per object ')
        sourcIDs_dic = collections.OrderedDict()
        for oo in np.arange(int(Nsources)):
            sourcIDs_dic[str(oo)] = [oo]
    else:
        sourcIDs_dic = source_association_dictionary
    Nobj = len(list(sourcIDs_dic.keys()))
    if verbose: print(' - Found '+str(Nobj)+' objects to generate spectra for in source_association_dictionary ')

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Assembling wavelength vector for spectra ')
    wavelengths     =  np.arange(model_cube_hdr['NAXIS3'])*model_cube_hdr['CD3_3']+model_cube_hdr['CRVAL3']
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    specfiles = []
    for oo, key in enumerate(sourcIDs_dic.keys()):
        obj_cube_hdr = model_cube_hdr.copy()
        try:
            specid       = str("%.10d" % int(key))
        except:
            specid       = str(key)
        specname     = outputdir+nameext+'_'+specid+'.fits'
        specfiles.append(specname)
        sourceIDs    = sourcIDs_dic[key]

        obj_cube_hdr.append(('OBJID   ',specid         ,'ID of object'),end=True)
        obj_cube_hdr.append(('SRCIDS  ',str(sourceIDs) ,'IDs of sources combined in object'),end=True)

        if verbose:
            infostr = ' - Extracting spectrum '+str("%6.f" % (oo+1))+' / '+str("%6.f" % Nobj)
            sys.stdout.write("%s\r" % infostr)
            sys.stdout.flush()

        sourceoutput = tes.extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=stddev_cube,
                                            source_model_cube=source_model_cube, data_cube=data_cube,
                                            specname=specname,obj_cube_hdr=obj_cube_hdr,clobber=clobber,verbose=True)

    if verbose: print('\n - Done extracting spectra. Returning list of fits files generated')
    return specfiles
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=None,source_model_cube=None,
                     specname='tdose_extract_spectra_extractedspec.fits',obj_cube_hdr=None,data_cube=None,
                     clobber=False,verbose=True):
    """
    Extracting a spectrum based on the layer scale image from the model cube provided a list of sources to combine.
    Noise is estimated from the noise cube (of the data)

    If all layer_scales are 1 a data_cube for the extractions is expected

    --- INPUT ---
    sourceIDs         The source IDs to combine into spectrum
    layer_scale_arr   Layer scale array (or image) produced when generating the model cube
                      fractional flux belonging to the source in each pixel
    wavelengths       Wavelength vector to use for extracted 1D spectrum.
    noise_cube        Cube with uncertainties (sqrt(variance)) of data cube to be used for estimating 1D uncertainties
                      To estimate S/N and 1D noise, providing a source model cube is required
    source_model_cube Source model cube containing the model cube for each individual source seperately
                      Needed in order to estimate noise from noise-cube
    specname          Name of file to save spectrum to
    obj_cube_hdr      Provide a template header to save the object cube (from combining the individual source cube)
                      as an extension to the extracted spectrum
    data_cube         In case all layers scales are 1, it is assumed that the source_model_cube contains a mask for the
                      spectral extraction, which will then be performed on this data_cube.
    clobber           To overwrite existing files set clobber=True
    verbose           Toggle verbosity

    --- EXAMPLE OF USE ---

    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Checking shape of wavelengths and layer_scale_arr')
    if wavelengths.shape[0] != layer_scale_arr.shape[1]:
        sys.exit(' ---> Shape of wavelength vector ('+str(wavelengths.shape)+
                 ') and wavelength dimension of layer scale array ('+
                 layer_scale_arr.shape[1].shape+') do not match.')
    else:
        if verbose: print('   dimensions match; proceeding...')
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Checking all sources have spectra in layer_scale_arr')
    maxsource = np.max(sourceIDs)
    if maxsource >= layer_scale_arr.shape[0]:
        sys.exit(' ---> Sources in list '+str(str(sourceIDs))+
                 ' not available among '+str(layer_scale_arr.shape[0])+' sources in layer_scale_arr.')
    else:
        if verbose: print('   All sources exist in layer_scale_arr; proceeding...')
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Assembling object spectrum from source scaling')
    source_ent = np.asarray(sourceIDs).astype(int)

    if (layer_scale_arr == 1).all():
        if verbose: print(' - All layer scales are 1; assuming source model cube contain mask for spectral extraction')
        object_cube  = np.sum(np.abs(source_model_cube[source_ent,:,:]),axis=0)
        if data_cube is None:
            sys.exit(' ---> Did not find a data cube to extrac spectra from as expected')

        object_mask     = (object_cube == 0) # masking all zeros in object mask
        invalid_mask    = np.ma.masked_invalid(data_cube).mask
        comb_mask       = (invalid_mask | object_mask)
        spec_1D_masked  = np.sum(np.sum(  np.ma.array(data_cube,mask=comb_mask)  ,axis=1),axis=1)
        spec_1D         = spec_1D_masked.filled(fill_value=0.0)

        if noise_cube is not None:
            if verbose: print('   Calculating noise as d_spec_k = sqrt( SUMij d_pix_ij**2 ), i.e., as the sqrt of variances summed')
            invalid_mask_noise = np.ma.masked_invalid(noise_cube).mask
            comb_mask          = (comb_mask | invalid_mask_noise)
            variance_1D_masked = np.ma.array(noise_cube,mask=comb_mask)**2
            noise_1D_masked    = np.sqrt( np.sum( np.sum( variance_1D_masked, axis=1), axis=1) )
            noise_1D           = noise_1D_masked.filled(fill_value=np.nan)

            if verbose: print('   Generating S/N vector')
            SN_1D         = spec_1D / noise_1D
        else:
            if verbose: print(' - No "noise_cube" provided. Setting all errors and S/N values to NaN')
            SN_1D    = np.zeros(spec_1D.shape)*np.NaN
            noise_1D = np.zeros(spec_1D.shape)*np.NaN
    else:
        if verbose: print(' - Some layer scales are different from 1; hence assembling spectra using layer scales')
        if len(source_ent) < 1:
            spec_1D   = layer_scale_arr[source_ent,:]
        else:
            spec_1D   = np.sum( layer_scale_arr[source_ent,:],axis=0)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        if noise_cube is not None:
            if verbose: print(' - Estimate S/N at each wavelength for 1D spectrum (see Eq. 16 of Kamann+2013)')
            if verbose: print('   Estimating fraction of flux in each pixel wrt. total flux in each layer')
            object_cube    = np.sum((source_model_cube[source_ent,:,:,:]),axis=0) # summing source models for all source IDs

            fluxfrac_cube_sents = np.zeros(source_model_cube.shape[1:])
            for sent in source_ent:
                object_cube_sent      = np.sum((source_model_cube[[sent],:,:,:]),axis=0) # getting source model for model 'sent'
                fluxscale1D_sent      = layer_scale_arr[sent,:]
                fluxfrac_cube_sent    = object_cube_sent / fluxscale1D_sent[:,None,None]
                fluxfrac_cube_sents   = fluxfrac_cube_sents + fluxfrac_cube_sent
            fluxfrac_cube = fluxfrac_cube_sents / len(source_ent) # renormalizing flux-fraction cube

            if verbose: print('   Defining pixel mask (ignoring NaN pixels) ') #+\
            #                  'and pixels with <'+str(fluxfrac_min)+' of total pixel flux in model cube) '
            # pix_mask      = (fluxfrac_cube < fluxfrac_min)
            invalid_mask1 = np.ma.masked_invalid(fluxfrac_cube).mask
            invalid_mask2 = np.ma.masked_invalid(noise_cube).mask

            # combining mask making sure all individual mask pixels have True for it to be true in combined mask
            comb_mask     = (invalid_mask1 | invalid_mask2) # | pix_mask

            if verbose: print('   Calculating noise propogated as d_spec_k = 1/sqrt( SUMij (fluxfrac_ij**2 / d_pix_ij**2) )')
            squared_ratio     = np.ma.array(fluxfrac_cube,mask=comb_mask)**2 / np.ma.array(noise_cube,mask=comb_mask)**2

            inv_noise_masked  = np.sqrt( np.sum( np.sum( squared_ratio, axis=1), axis=1) )
            noise_1D          = (1.0/inv_noise_masked).filled(fill_value=0.0)
            if verbose: print('   Generating S/N vector')
            SN_1D         = spec_1D / noise_1D
        else:
            if verbose: print(' - No "noise_cube" provided. Setting all errors and S/N values to NaN')
            SN_1D    = np.zeros(spec_1D.shape)*np.NaN
            noise_1D = np.zeros(spec_1D.shape)*np.NaN

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Saving extracted 1D spectrum and source cube to \n   '+specname)
    mainHDU = afits.PrimaryHDU()       # primary HDU
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    c1 = afits.Column(name='wave',      format='D', unit='ANGSTROMS', array=wavelengths)
    c2 = afits.Column(name='flux',      format='D', unit='', array=spec_1D)
    c3 = afits.Column(name='fluxerror', format='D', unit='', array=noise_1D)
    c4 = afits.Column(name='s2n',       format='D', unit='', array=SN_1D)

    coldefs = afits.ColDefs([c1,c2,c3,c4])
    th = afits.BinTableHDU.from_columns(coldefs) # creating default header

    # writing hdrkeys:'---KEY--',                             '----------------MAX LENGTH COMMENT-------------'
    th.header.append(('EXTNAME ','SPEC1D'                     ,'cube containing source'),end=True)
    head    = th.header

    tbHDU  = afits.BinTableHDU.from_columns(coldefs, header=head)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if obj_cube_hdr is not None:
        objHDU        = afits.ImageHDU(object_cube)
        for hdrkey in list(obj_cube_hdr.keys()):
            if not hdrkey in list(objHDU.header.keys()):
                objHDU.header.append((hdrkey,obj_cube_hdr[hdrkey],obj_cube_hdr.comments[hdrkey]),end=True)

        try:
            objHDU.header.append(('EXTNAMEC',objHDU.header['EXTNAME'] ,'EXTNAME of original source cube'),end=True)
            del objHDU.header['EXTNAME']
        except:
            pass

        objHDU.header.append(('EXTNAME ','SOURCECUBE'            ,'cube containing source'),end=True)

        hdulist = afits.HDUList([mainHDU,tbHDU,objHDU])
    else:
        hdulist = afits.HDUList([mainHDU,tbHDU])
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    hdulist.writeto(specname, overwrite=clobber)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    return wavelengths, spec_1D, noise_1D, object_cube
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectra_viasourcemodelcube(datacube,sourcemodelcube,wavelengths,speclist,specids='None',outputdir='./',
                                       noisecube=False,sourcemodel_hdr='None',verbose=True):
    """
    Wrapper for tes.extract_spectrum_viasourcemodelcube() to extract mutliple spectra

    --- INPUT ----
    datacube          Datacube to extract spectra from
    sourcemodelcube   Cube containing the source models for each object used as "extraction cube"
                      Dimensions should be [Nsources,datacube.shape]
    wavelengths       Wavelength vector to use for extracted 1D spectrum.
    speclist          List of spectra to extract. Indexes corresponding to the source models in the
                      sourcemodlecube
    specids           List of IDs to use in naming of output for source models referred to in "speclist"
    outputdir         Directory to store spectra to
    noisecube         Cube with uncertainties (sqrt(variance)) of data cube to be used in extraction
    souremodel_hdr    If not 'None' provide a basic fits header for the source model cubes extracted
                      and they will be appended to the individual output fits file containing the extracted
                      spectra.
    verbose           Toggle verbosity

    --- EXAMPLE OF USE ---

    """
    if verbose: print(' - Check that source models indicated are present in source model cube ')
    specnames  = []
    Nmodels    = sourcemodelcube.shape[0]
    maxobj     = np.max(speclist)
    if maxobj >= Nmodels:
        sys.exit(' ---> Object model "'+str(maxobj)+'" is not included in source model cube (models start at 0)')
    else:
        if verbose: print('   All object models appear to be included in the '+str(Nmodels)+' source models found in cube')

    if datacube.shape != sourcemodelcube[0].shape:
        sys.exit(' ---> Shape of datacube ('+str(datacube.shape)+') and shape of source models ('+
                 sourcemodelcube[0].shape+') do not match.')

    sourcemodel_sum = np.sum(sourcemodelcube,axis=0)
    for ss, spec in enumerate(speclist):
        if specids == 'None':
            specid = spec
        else:
            specid = specids[ss]

        specname = outputdir+'tdose_spectrum_'+str("%.12d" % specid)+'.fits'
        specnames.append(specname)

        sourcemodel     = sourcemodelcube[spec,:,:,:]
        sourceweights   = sourcemodel/sourcemodel_sum  # fractional flux of model for given source in each pixel
        sourcemodel_hdr.append(('OBJMODEL',spec      ,'Source model number in parent source model cube'),end=True)
        sourcemodel_hdr.append(('OBJID   ',specid    ,'ID of source'),end=True)

        if verbose:
            infostr = ' - Extracting spectrum '+str("%6.f" % (spec+1))+' / '+str("%6.f" % len(speclist))
            sys.stdout.write("%s\r" % infostr)
            sys.stdout.flush()

        sourceoutput = tes.extract_spectrum_viasourcemodelcube(datacube,sourceweights,wavelengths,specname=specname,
                                                               noisecube=noisecube,spec1Dmethod='sum',
                                                               sourcecube_hdr=sourcemodel_hdr,verbose=verbose)

    if verbose: print('\n - Done extracting spectra. Returning list of fits files containing spectra')
    return specnames
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectrum_viasourcemodelcube(datacube,sourceweights,wavelengths,
                                        specname='tdose_extract_spectra_extractedspec.fits',
                                        noisecube=None,spec1Dmethod='sum',sourcecube_hdr='None',verbose=True):
    """
    Extracting a spectrum from a data cube given a source model (cube) to be used as 'extraction cube'

    --- INPUT ---
    datacube          Datacube to extract spectra from
    sourceweights     Weights from source model to use as "extraction cube". The weights should contain the
                      fractional flux belonging to the source in each pixel
    wavelengths       Wavelength vector to use for extracted 1D spectrum.
    specname          Name of spectrum to generate
    noisecube         Cube with uncertainties (sqrt(variance)) of data cube to be used in extraction
    spec1Dmethod      Method used to extract 1D spectrum from source cube with
    sourcecube_hdr    If not 'None' provide a fits header for the source cube and it ill be appended to the
                      output fits file.
    verbose           Toggle verbosity

    --- EXAMPLE OF USE ---

    """
    if verbose: print(' - Checking shape of data and source model cubes')

    if datacube.shape != sourceweights.shape:
        sys.exit(' ---> Shape of datacube ('+str(datacube.shape)+') and source weights ('+
                 sourceweights.shape+') do not match.')
    else:
        if verbose: print('   dimensions match; proceeding with extraction ')
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Applying weights to "datacube" to obtain source cube ')
    sourcecube     = datacube*sourceweights

    if noisecube is not None:
        if verbose: print(' - Using "noisecube" for error propagation ')
        datanoise = noisecube
    else:
        if verbose: print(' - No "noisecube" provided. Setting all errors to 1')
        datanoise = np.ones(datacube.shape)

    if verbose: print(' - Assuming uncertainty on source weights equals the datanoise when propgating errors')
    sourceweights_err = datanoise
    sourcecube_err    = sourcecube * np.sqrt( (datanoise/datacube)**2 + (sourceweights_err/sourceweights)**2 )

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Generating 1D spectrum from source cube via:')
    spec_wave   = wavelengths
    maskinvalid = np.ma.masked_invalid(sourcecube * sourcecube_err).mask
    if spec1Dmethod == 'sum':
        if verbose: print('   Simple summation of fluxes in sourcecube.')
        spec_flux = np.sum(np.sum(np.ma.array(sourcecube,mask=maskinvalid),axis=1),axis=1).filled()
        if verbose: print('   Errors are propagated as sum of squares.')
        spec_err  = np.sqrt( np.sum( np.sum(np.ma.array(sourcecube_err,mask=maskinvalid)**2,axis=1),axis=1) ).filled()
    elif spec1Dmethod == 'sum_SNweight':
        pdb.set_trace()
    else:
        sys.exit(' ---> The chosen spec1Dmethod ('+str(spec1Dmethod)+') is invalid')

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print(' - Saving extracted 1D spectrum and source cube to \n   '+specname)
    mainHDU = afits.PrimaryHDU()       # primary HDU
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    c1 = afits.Column(name='wave', format='D', unit='ANGSTROMS', array=spec_wave)
    c2 = afits.Column(name='flux', format='D', unit='', array=spec_flux)
    c3 = afits.Column(name='fluxerror', format='D', unit='', array=spec_err)

    coldefs = afits.ColDefs([c1,c2,c3])
    th = afits.BinTableHDU.from_columns(coldefs) # creating default header

    # writing hdrkeys:'---KEY--',                             '----------------MAX LENGTH COMMENT-------------'
    th.header.append(('EXTNAME ','SPEC1D'                     ,'cube containing source'),end=True)
    th.header.append(('SPECMETH' , spec1Dmethod               ,'Method used for spectral extraction'),end=True)
    head    = th.header

    tbHDU  = afits.BinTableHDU.from_columns(coldefs, header=head)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if sourcecube_hdr != 'None':
        sourceHDU        = afits.ImageHDU(sourcecube)       # default HDU with default minimal header
        for hdrkey in list(sourcecube_hdr.keys()):
            if not hdrkey in list(sourceHDU.header.keys()):
                sourceHDU.header.append((hdrkey,sourcecube_hdr[hdrkey],sourcecube_hdr.comments[hdrkey]),end=True)

        sourceHDU.header.append(('EXTNAME ','SOURCECUBE'            ,'cube containing source'),end=True)

        hdulist = afits.HDUList([mainHDU,tbHDU,sourceHDU])
    else:
        hdulist = afits.HDUList([mainHDU,tbHDU])
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    hdulist.writeto(specname, overwrite=True)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    return sourcecube, sourcecube_err, spec_wave, spec_flux, spec_err
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def plot_1Dspecs(filelist,plotname='./tdose_1Dspectra.pdf',colors=None,labels=None,plotSNcurve=False,
                 tdose_wavecol='wave',tdose_fluxcol='flux',tdose_errcol='fluxerror',
                 simsources=None,simsourcefile='/Users/kschmidt/work/TDOSE/mock_cube_sourcecat161213_all.fits',
                 sim_cube_dim=None,comparisonspecs=None,comp_colors=['blue'],comp_labels=None,
                 comp_wavecol='WAVE_AIR',comp_fluxcol='FLUX',comp_errcol='FLUXERR',
                 xrange=None,yrange=None,showspecs=False,shownoise=True,
                 skyspecs=None,sky_colors=['red'],sky_labels=['sky'],
                 sky_wavecol='lambda',sky_fluxcol='data',sky_errcol='stat',
                 showlinelists=None,linelistcolors=['gray'],smooth=0,ylog=False,
                 plotratio=False,
                 verbose=True,pubversion=False):
    """
    Plots of multiple 1D spectra

    --- INPUT ---
    filelist            List of spectra filenames to plot
    plotname            Name of plot to generate
    colors              Colors of the spectra in filelist to use
    labels              Labels of the spectra in filelist to use
    plotSNcurve         Show signal-to-noise curve instead of flux spectra
    tdose_wavecol       Wavelength column of the spectra in filelist
    tdose_fluxcol       Flux column of the spectra in filelist
    tdose_errcol        Flux error column of the spectra in filelist
    simsources          To plot simulated sources provide ids here
    simsourcefile       Source file with simulated sources to plot
    sim_cube_dim        Dimensions of simulated cubes
    comparisonspecs     To plot comparison spectra provide the filenames of those here
    comp_colors         Colors of the spectra in comparisonspecs list to use
    comp_labels         Labels of the spectra in comparisonspecs list to use
    comp_wavecol        Wavelength column of the spectra in comparisonspecs list
    comp_fluxcol        Flux column of the spectra in comparisonspecs list
    comp_errcol         Flux error column of the spectra in comparisonspecs list
    xrange              Xrange of plot
    yrange              Yrange of plot
    showspecs           To show plot instead of storing it to disk set showspecs=True
    shownoise           To add noise envelope around spectrum set shownoise=True
    skyspecs            To plot sky spectra provide the filenames of those here
    sky_colors          Colors of the spectra in skyspecs list to use
    sky_labels          Labels of the spectra in skyspecs list to use
    sky_wavecol         Wavelength column of the spectra in skyspecs list
    sky_fluxcol         Flux column of the spectra in skyspecs list
    sky_errcol          Flux error column of the spectra in skyspecs list
    showlinelists       To show line lists provide a list of arrays of dimension (Nlines,2) where each row in the
                        arrays contains [waveobs, name], where 'waveobs' is the observed wavelengths and 'name' is
                        a string with the name of each of the Nlines postions to mark on the spectrum.
    linelistcolors      List of colors for line lists provided in showlinelists
    smooth              To smooth the spectra, provide sigma of the 1D gaussian smoothing kernel to apply.
                        For smooth = 0, no smoothing is performed.
    ylog                To plot y-axis in log scale set to true
    plotratio           To plot the ratio between the main spectrum and the comparison spectra instead of the actual
                        spectra, set this keyword to true.
    verbose             Toggle verbosity
    pubversion          Generate more publication friendly version of figure

    """
    if len(filelist) == 1:
        if verbose: print(' - Plotting data from '+filelist[0])
    else:
        if verbose: print(' - Plotting data from filelist ')

    if pubversion:
        fig = plt.figure(figsize=(6, 3))
        fig.subplots_adjust(wspace=0.1, hspace=0.1,left=0.15, right=0.95, bottom=0.18, top=0.83)
        Fsize  = 12
    else:
        fig = plt.figure(figsize=(10, 3))
        fig.subplots_adjust(wspace=0.1, hspace=0.1,left=0.06, right=0.81, bottom=0.15, top=0.95)
        Fsize  = 10

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # Looking for flux units in spectra
    bunit = 'BUNIT FLUX' # Default BUNIT
    for unitspec in filelist:

        if bunit == 'BUNIT FLUX':
            try:
                sourcecubehdr = afits.open(unitspec)['SOURCECUBE'].header
                bunit         = sourcecubehdr['BUNIT']
            except:
                try: # Backwards compatibility to TDOSE v2.0 extractions
                    sourcecubehdr = afits.open(unitspec)[2].header
                    bunit         = sourcecubehdr['BUNIT']
                except:
                    pass
    if bunit == 'BUNIT FLUX':
        if verbose: print(' - Did not find BUNIT in SOURCECUBE header for any spectra in filelist - are they not from TDOSE?')

    if bunit == '10**(-20)*erg/s/cm**2/Angstrom': # Making bunit LaTeXy for MUSE-Wide BUNIT format
        bunit = '1e-20 erg/s/cm$^2$/\AA'
    else:
        bunit = '$'+bunit+'$' # minimizing pronlems with LaTeXing plot axes
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    lthick = 1
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif',size=Fsize)
    plt.rc('xtick', labelsize=Fsize)
    plt.rc('ytick', labelsize=Fsize)
    plt.clf()
    plt.ioff()
    #plt.title(plotname.split('TDOSE 1D spectra'),fontsize=Fsize)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    for ff, specfile in enumerate(filelist):
        specdat = afits.open(specfile)[1].data

        if colors is None:
            spec_color = None
        else:
            spec_color = colors[ff]

        if labels is None:
            spec_label = specfile
        else:
            spec_label = labels[ff]

        if xrange is not None:
            goodent = np.where((specdat[tdose_wavecol] > xrange[0]) & (specdat[tdose_wavecol] < xrange[1]))[0]
            if goodent == []:
                if verbose: print(' - The chosen xrange is not covered by the input spectrum. Plotting full spectrum')
                goodent = np.arange(len(specdat[tdose_wavecol]))
        else:
            goodent = np.arange(len(specdat[tdose_wavecol]))

        if plotSNcurve:
            try:
                s2ndat = specdat['s2n'][goodent]
            except:
                s2ndat = specdat[tdose_fluxcol][goodent]/specdat[tdose_errcol][goodent]

            if smooth > 0:
                s2ndat = snf.gaussian_filter(s2ndat, smooth)

            if not plotratio:
                plt.plot(specdat[tdose_wavecol][goodent],s2ndat,color=spec_color,lw=lthick, label=spec_label)
                ylabel = 'S/N'
            else:
                plt.plot(specdat[tdose_wavecol][goodent],s2ndat/s2ndat,color=spec_color,lw=lthick, label=None)
                ylabel = 'S/N ratio'
            #plotname = plotname.replace('.pdf','_S2N.pdf')
        else:
            fillalpha = 0.30
            fluxdat   = specdat[tdose_fluxcol][goodent]
            errlow    = specdat[tdose_fluxcol][goodent]-specdat[tdose_errcol][goodent]
            errhigh   = specdat[tdose_fluxcol][goodent]+specdat[tdose_errcol][goodent]

            if smooth > 0:
                fluxdat = snf.gaussian_filter(fluxdat, smooth)
                errlow      = snf.gaussian_filter(errlow,  smooth)
                errhigh     = snf.gaussian_filter(errhigh, smooth)

            if smooth > 0:
                fluxdat = snf.gaussian_filter(fluxdat, smooth)

            if not plotratio:
                if shownoise:
                    plt.fill_between(specdat[tdose_wavecol][goodent],errlow,errhigh,
                                     alpha=fillalpha,color=spec_color)

                plt.plot(specdat[tdose_wavecol][goodent],fluxdat,
                         color=spec_color,lw=lthick, label=spec_label)
                ylabel = tdose_fluxcol
            else:
                plt.plot(specdat[tdose_wavecol][goodent],fluxdat/fluxdat,
                         color=spec_color,lw=lthick, label=None)
                ylabel = tdose_fluxcol+' ratio '
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if simsources is not None:
        sim_total = np.zeros(len(specdat[tdose_wavecol]))
        for sourcenumber in simsources:
            sourcedat = afits.open(simsourcefile)[1].data
            xpos       = sourcedat['xpos'][sourcenumber]
            ypos       = sourcedat['ypos'][sourcenumber]
            fluxscale  = sourcedat['fluxscale'][sourcenumber]
            sourcetype = sourcedat['sourcetype'][sourcenumber]
            spectype   = sourcedat['spectype'][sourcenumber]
            sourcecube = tbmc.gen_source_cube([ypos,xpos],fluxscale,sourcetype,spectype,cube_dim=sim_cube_dim,
                                              verbose=verbose,showsourceimgs=False)

            simspec    = np.sum( np.sum(sourcecube, axis=1), axis=1)
            sim_total  = sim_total + simspec
            if smooth > 0:
                simspec = snf.gaussian_filter(simspec, smooth)

            plt.plot(specdat[tdose_wavecol],simspec,'--',color='black',lw=lthick)

        plt.plot(specdat[tdose_wavecol],sim_total,'--',color='black',lw=lthick,
                 label='Sim. spectrum: \nsimsource='+str(simsources))

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if comparisonspecs is not None:
        for cc, comparisonspec in enumerate(comparisonspecs):
            compdat = afits.open(comparisonspec)[1].data

            if xrange is not None:
                goodent = np.where((compdat[comp_wavecol] > xrange[0]) & (compdat[comp_wavecol] < xrange[1]))[0]
                if goodent == []:
                    if verbose: print(' - The chosen xrange is not covered by the comparison spectrum. Plotting full spectrum')
                    goodent = np.arange(len(compdat[comp_wavecol]))
            else:
                goodent = np.arange(len(compdat[comp_wavecol]))

            if comp_colors is None:
                comp_color = None
            else:
                comp_color = comp_colors[cc]

            if comp_labels is None:
                comp_label = comparisonspec
            else:
                comp_label = comp_labels[cc]

            if plotSNcurve:
                s2ncompdat = compdat[comp_fluxcol][goodent]/compdat[comp_errcol][goodent]
                if smooth > 0:
                    s2ncompdat = snf.gaussian_filter(s2ncompdat, smooth)

                if not plotratio:
                    plt.plot(compdat[comp_wavecol][goodent],s2ncompdat,
                             color=comp_color,lw=lthick, label=comp_label)
                else:
                    plt.plot(compdat[comp_wavecol][goodent],s2ndat/s2ncompdat,
                             color=comp_color,lw=lthick, label=comp_label)

            else:
                fillalpha = 0.30
                fluxcompdat = compdat[comp_fluxcol][goodent]
                errlow      = compdat[comp_fluxcol][goodent]-compdat[comp_errcol][goodent]
                errhigh     = compdat[comp_fluxcol][goodent]+compdat[comp_errcol][goodent]

                if smooth > 0:
                    fluxcompdat = snf.gaussian_filter(fluxcompdat, smooth)
                    errlow      = snf.gaussian_filter(errlow,  smooth)
                    errhigh     = snf.gaussian_filter(errhigh, smooth)

                if not plotratio:
                    if shownoise:
                        plt.fill_between(compdat[comp_wavecol][goodent],errlow,errhigh,
                                         alpha=fillalpha,color=comp_color)

                    plt.plot(compdat[comp_wavecol][goodent],fluxcompdat,
                             color=comp_color,lw=lthick, label=comp_label)
                else:
                    plt.plot(compdat[comp_wavecol][goodent],fluxdat/fluxcompdat,
                             color=comp_color,lw=lthick, label=comp_label)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if skyspecs is not None:
        for ss, skyspec in enumerate(skyspecs):
            skydat = afits.open(skyspec)[1].data

            if xrange is not None:
                goodent = np.where((skydat[sky_wavecol] > xrange[0]) & (skydat[sky_wavecol] < xrange[1]))[0]
                if goodent == []:
                    if verbose: print(' - The chosen xrange is not covered by the sky spectrum. Plotting full spectrum')
                    goodent = np.arange(len(skydat[sky_wavecol]))
            else:
                goodent = np.arange(len(skydat[sky_wavecol]))

            if sky_colors is None:
                sky_color = None
            else:
                sky_color = sky_colors[ss]

            if sky_labels is None:
                sky_label = skyspec
            else:
                sky_label = sky_labels[ss]

            if plotSNcurve:
                s2nsky = skydat[sky_fluxcol][goodent]/skydat[sky_errcol][goodent]
                if smooth > 0:
                    s2nsky = snf.gaussian_filter(s2nsky, smooth)

                plt.plot(skydat[sky_wavecol][goodent],s2nsky,
                         color=sky_color,lw=lthick, label=sky_label)
            else:
                fillalpha = 0.30
                fluxsky   = skydat[sky_fluxcol][goodent]
                errlow    = skydat[sky_fluxcol][goodent]-skydat[sky_errcol][goodent]
                errhigh   = skydat[sky_fluxcol][goodent]+skydat[sky_errcol][goodent]

                if smooth > 0:
                    fluxsky = snf.gaussian_filter(fluxsky, smooth)
                    errlow  = snf.gaussian_filter(errlow,  smooth)
                    errhigh = snf.gaussian_filter(errhigh, smooth)

                if shownoise:
                    plt.fill_between(skydat[sky_wavecol][goodent],errlow,errhigh,
                                     alpha=fillalpha,color=sky_color)

                plt.plot(skydat[sky_wavecol][goodent],fluxsky,
                         color=sky_color,lw=lthick, label=sky_label)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if xrange is None:
        xvals = [4800,9300]
    else:
        xvals = xrange
    plt.plot(xvals,[0,0],'--k',lw=lthick)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    plt.xlabel('Wavelength [\AA]', fontsize=Fsize)

    if pubversion:
        if plotSNcurve:
            ylabel = 'Signal-to-Noise'
        else:
            ylabel = 'Flux ['+str(bunit)+']'

        if plotratio:
            ylabel = ylabel+' ratio'

    plt.ylabel(ylabel, fontsize=Fsize)

    if ylog:
        plt.yscale('log')

    if yrange is not None:
        plt.ylim(yrange)

    if xrange is not None:
        plt.xlim(xrange)

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if showlinelists is not None:
        for sl, showlinelist in enumerate(showlinelists):
            ymin, ymax = plt.ylim()
            xmin, xmax = plt.xlim()
            for ww, wave in enumerate(showlinelist[:,0]):
                wave = float(wave)
                if (wave < xmax) & (wave > xmin):
                    plt.plot([wave,wave],[ymin,ymax],linestyle='--',color=linelistcolors[sl],lw=lthick)
                    plt.text(wave,ymin+1.03*np.abs([ymax-ymin]),showlinelist[:,1][ww],color=linelistcolors[sl], fontsize=Fsize-2., ha='center')
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if pubversion:
        leg = plt.legend(fancybox=True, loc='upper center',prop={'size':Fsize-2},ncol=4,numpoints=1,
                         bbox_to_anchor=(0.44, 1.27))  # add the legend
    else:
        leg = plt.legend(fancybox=True, loc='upper right',prop={'size':Fsize},ncol=1,numpoints=1,
                         bbox_to_anchor=(1.25, 1.03))  # add the legend

    leg.get_frame().set_alpha(0.7)

    if showspecs:
        if verbose: print('   Showing plot (not saving to file)')
        plt.show()
    else:
        if verbose: print('   Saving plot to',plotname)
        plt.savefig(plotname)

    plt.clf()
    plt.close('all')
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def plot_histograms(datavectors,plotname='./tdose_cubehist.pdf',colors=None,labels=None,bins=None,
                    xrange=None,yrange=None,verbose=True,norm=True,ylog=True):
    """
    Plot histograms of a set of data vectors.

    --- INPUT ---
    datavectors     Set of data vectors to plot histograms of
    plotname        Name of plot to generate
    colors          Colors to use for histograms
    labels          Labels for the data vectors
    bins            Bins to use for histograms. Can be generated with np.arange(minval,maxval+binwidth,binwidth)
    xrange          Xrange of plot
    yrange          Yrange of plot
    verbose         Toggle verbosity
    norm            Noramlize the histograms
    ylog            Use a logarithmic y-axes when plotting

    """
    Ndat = len(datavectors)
    if verbose: print(' - Plotting histograms of N = '+str(Ndat)+' data vectors')

    if colors is None:
        colors = ['blue']*Ndat

    if labels is None:
        labels = ['data vector no. '+str(ii+1) for ii in np.arange(Ndat)]

    if bins is None:
        bins = np.arange(-100,102,2)

    fig = plt.figure(figsize=(10, 3))
    fig.subplots_adjust(wspace=0.1, hspace=0.1,left=0.08, right=0.81, bottom=0.1, top=0.95)
    Fsize  = 10
    lthick = 1
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif',size=Fsize)
    plt.rc('xtick', labelsize=Fsize)
    plt.rc('ytick', labelsize=Fsize)
    plt.clf()
    plt.ioff()
    #plt.title(plotname.split('TDOSE 1D spectra'),fontsize=Fsize)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    for dd, datavec in enumerate(datavectors):
        hist = plt.hist(datavec[~np.isnan(datavec)],color=colors[dd],bins=bins,histtype="step",lw=lthick,
                        label=labels[dd],normed=norm)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if yrange is None:
        yvals = [1e-5,1e8]
    else:
        yvals = yrange
    plt.plot([0,0],yvals,'--k',lw=lthick)
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    plt.xlabel('', fontsize=Fsize)
    plt.ylabel('\#', fontsize=Fsize)

    if yrange is not None:
        plt.ylim(yrange)

    if xrange is not None:
        plt.xlim(xrange)

    if ylog:
        plt.yscale('log')

    leg = plt.legend(fancybox=True, loc='upper right',prop={'size':Fsize},ncol=1,numpoints=1,
                     bbox_to_anchor=(1.25, 1.03))  # add the legend
    leg.get_frame().set_alpha(0.7)

    if verbose: print('   Saving plot to',plotname)
    plt.savefig(plotname)
    plt.clf()
    plt.close('all')
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =


""" Factory methods to create N-D panels """

import warnings
from pandas.compat import zip
import pandas.compat as compat

def create_nd_panel_factory(klass_name, orders, slices, slicer, aliases=None,
                            stat_axis=2, info_axis=0, ns=None):
    """ manufacture a n-d class:

    DEPRECATED. Panelnd is deprecated and will be removed in a future version.
    The recommended way to represent these types of n-dimensional data are with
    the `xarray package <http://xarray.pydata.org/en/stable/>`__.
    Pandas provides a `.to_xarray()` method to automate this conversion.

    Parameters
    ----------
    klass_name : the klass name
    orders : the names of the axes in order (highest to lowest)
    slices : a dictionary that defines how the axes map to the slice axis
    slicer : the class representing a slice of this panel
    aliases : a dictionary defining aliases for various axes
        default = { major : major_axis, minor : minor_axis }
    stat_axis : the default statistic axis default = 2
    info_axis : the info axis

    Returns
    -------
    a class object representing this panel
    """

    # if slicer is a name, get the object
    if isinstance(slicer, compat.string_types):
        import pandas
        try:
            slicer = getattr(pandas, slicer)
        except:
            raise Exception("cannot create this slicer [%s]" % slicer)

    # build the klass
    ns = {} if not ns else ns
    klass = type(klass_name, (slicer, ), ns)

    # setup the axes
    klass._setup_axes(axes=orders, info_axis=info_axis, stat_axis=stat_axis,
                      aliases=aliases, slicers=slices)

    klass._constructor_sliced = slicer

    # define the methods ####
    def __init__(self, *args, **kwargs):

        # deprecation GH13564
        warnings.warn("\n{klass} is deprecated and will be removed in a "
                      "future version.\nThe recommended way to represent "
                      "these types of n-dimensional data are with the\n"
                      "`xarray package "
                      "<http://xarray.pydata.org/en/stable/>`__.\n"
                      "Pandas provides a `.to_xarray()` method to help "
                      "automate this conversion.\n".format(
                          klass=self.__class__.__name__),
                      FutureWarning, stacklevel=2)

        if not (kwargs.get('data') or len(args)):
            raise Exception("must supply at least a data argument to [%s]" %
                            klass_name)
        if 'copy' not in kwargs:
            kwargs['copy'] = False
        if 'dtype' not in kwargs:
            kwargs['dtype'] = None
        self._init_data(*args, **kwargs)

    klass.__init__ = __init__

    def _get_plane_axes_index(self, axis):
        """ return the sliced index for this object """

        # TODO: axis_name is not used, remove?
        axis_name = self._get_axis_name(axis)  # noqa
        index = self._AXIS_ORDERS.index(axis)

        planes = []
        if index:
            planes.extend(self._AXIS_ORDERS[0:index])
        if index != self._AXIS_LEN:
            planes.extend(self._AXIS_ORDERS[index + 1:])

        return planes

    klass._get_plane_axes_index = _get_plane_axes_index

    def _combine(self, other, func, axis=0):
        if isinstance(other, klass):
            return self._combine_with_constructor(other, func)
        return super(klass, self)._combine(other, func, axis=axis)

    klass._combine = _combine

    def _combine_with_constructor(self, other, func):

        # combine labels to form new axes
        new_axes = []
        for a in self._AXIS_ORDERS:
            new_axes.append(getattr(self, a).union(getattr(other, a)))

        # reindex: could check that everything's the same size, but forget it
        d = dict([(a, ax) for a, ax in zip(self._AXIS_ORDERS, new_axes)])
        d['copy'] = False
        this = self.reindex(**d)
        other = other.reindex(**d)

        result_values = func(this.values, other.values)

        return self._constructor(result_values, **d)

    klass._combine_with_constructor = _combine_with_constructor

    # set as NonImplemented operations which we don't support
    for f in ['to_frame', 'to_excel', 'to_sparse', 'groupby', 'join', 'filter',
              'dropna', 'shift']:

        def func(self, *args, **kwargs):
            raise NotImplementedError("this operation is not supported")

        setattr(klass, f, func)

    # add the aggregate operations
    klass._add_aggregate_operations()
    klass._add_numeric_operations()

    return klass


#!/usr/bin/env python
"""
=================================
Path with L1- Logistic Regression
=================================

Computes path on IRIS dataset.

"""
print(__doc__)

# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD 3 clause

from datetime import datetime
import numpy as np
import matplotlib.pyplot as plt

from sklearn import linear_model
from sklearn import datasets
from sklearn.svm import l1_min_c

iris = datasets.load_iris()
X = iris.data
y = iris.target

X = X[y != 2]
y = y[y != 2]

X -= np.mean(X, 0)

###############################################################################
# Demo path functions

cs = l1_min_c(X, y, loss='log') * np.logspace(0, 3)

print("Computing regularization path ...")
start = datetime.now()
clf = linear_model.LogisticRegression(C=1.0, penalty='l1', tol=1e-6)
coefs_ = []
for c in cs:
    clf.set_params(C=c)
    clf.fit(X, y)
    coefs_.append(clf.coef_.ravel().copy())
print("This took ", datetime.now() - start)

coefs_ = np.array(coefs_)
plt.plot(np.log10(cs), coefs_)
ymin, ymax = plt.ylim()
plt.xlabel('log(C)')
plt.ylabel('Coefficients')
plt.title('Logistic Regression Path')
plt.axis('tight')
plt.show()


#!/usr/bin/env python
#-*- coding:utf-8 -*-

import matplotlib
from matplotlib import pyplot as plt

import numpy as np
from numpy import linalg as LA
import cv2
import math

import plotly
from plotly.graph_objs import *
from plotly.offline import download_plotlyjs, init_notebook_mode, iplot
from plotly import tools

import time
import collections as col
from collections import OrderedDict

import ast
from ndreg import *
import ndio.remote.neurodata as neurodata

import nibabel as nib
import networkx as nx
import re
import pandas as pd
import requests
import json
import seaborn as sns
import csv, gc

from sklearn.manifold import spectral_embedding as se

import scipy.sparse as sp

plotly.offline.init_notebook_mode()

def spec_clust(graphx, num_components):
    """
    Function for doing the spectral embedding.
    :param graphx:
    :param num_components:
    :return:
    """
    adj_mat = nx.adjacency_matrix(graphx)
#     result = se(adj_mat, n_components=num_components, drop_first=True)
    result = se(adj_mat, n_components=num_components, drop_first=False)

    return result

def add_to_dict(d, region, index):
    if region in d:
        d[region].append(index)
    else:
        d[region] = [index]

    return d

def get_adj_mat(regions_path):
    points = np.genfromtxt(regions_path, delimiter=',')
    x_dim = np.max(points[:, 0])
    y_dim = np.max(points[:, 1])
    z_dim = np.max(points[:, 2])
    am = SparseMatrix(x_dim, y_dim, z_dim)
    for point in points:
        am.add(tuple(point[0:3]), point[4])

    return am

def get_dict_real(g, se_result, regions_path):
    nodes = g.nodes()
    points = np.genfromtxt(regions_path, delimiter=',')
    orig_dict = OrderedDict()
    d = {}

    sparse_mat = get_adj_mat(regions_path)

    for index, node in enumerate(nodes):
        s = g.node[node]['attr']
        point = ast.literal_eval(s)
        region = sparse_mat.get(tuple(point))
        #         if region == -1:
        #             # error
        #             print 'FUCK'
        add_to_dict(d, region, index)

    for point in points:
        region = point[4]
        #         if region in orig_dict:
        #             orig_dict[region] = np.vstack((orig_dict[region], point[0:3]))
        #         else:
        #             orig_dict[region] = np.array([point[0:3]])
        add_to_dict(orig_dict, region, point[0:3])

    se_regions_nodes = {}
    se_regions = {}
    for key, value in d.iteritems():
        index_list = value
        nodes_arr = np.array(nodes)
        pt_names = nodes_arr[index_list]
        se_pts = se_result[index_list]
        nodes_to_se = dict(zip(pt_names, se_pts))  # maps from node names to embedded point coordinates
        se_regions_nodes[key] = nodes_to_se
        se_regions[key] = se_pts

    return se_regions, orig_dict, se_regions_nodes

def create_connectivity_graph(orig_avg_dict, se_avg_dict, max_dist=0.02):
    g = nx.Graph()
    for key, avg in se_avg_dict.iteritems():
        for key2, avg2 in se_avg_dict.iteritems():
            avg_np = np.array(avg)
            avg2_np = np.array(avg2)
            diff = np.linalg.norm(avg_np - avg2_np)
            diff = max_dist if diff > max_dist else diff
            g.add_edge(key, key2, weight=diff)

    # Setting the coordinate attribute for each region node to the average of that region.
    for key, avg in orig_avg_dict.iteritems():
        g.node[key]['attr'] = avg

    return g

def get_connectivity_hard(eig_dict, orig_dict=None, max_dist=0.02):
    """
    Uses create_connectivity_graph.
    :param eig_dict:
    :param orig_dict:
    :return:
    """
    eigenvector_index = 1  # the second smallest eigenvector
    avg_dict = {}
    orig_avg_dict = OrderedDict()

    # dict that maps from region to most connected region
    con_dict = OrderedDict()

    orig_con_dict = OrderedDict()

    if orig_dict != None:
        # Getting the original averages.
        for key, region in orig_dict.iteritems():
            tmp_x = []
            tmp_y = []
            y_vals = []

            for j in range(len(region)):
                y_vals.append(region[j])
            y_vals = np.array(y_vals)
            x_avg = np.mean(y_vals[:, 0])
            y_avg = np.mean(y_vals[:, 1])
            z_avg = np.mean(y_vals[:, 2])
            orig_avg_dict[key] = [x_avg, y_avg, z_avg]
        # avg = np.mean(y_vals)
        #             orig_avg_dict[key] = avg

        #         print 'orignal averages'
        #         print orig_avg_dict

        # Getting connectivity for original points.
        for key, avg in orig_avg_dict.iteritems():
            min_key = ''
            min_diff = float('inf')
            for key2, avg2 in orig_avg_dict.iteritems():
                if key2 == key:
                    continue
                avg_np = np.array(avg)
                avg2_np = np.array(avg2)
                diff = np.linalg.norm(avg_np - avg2_np)
                if diff < min_diff:
                    min_diff = diff
                    min_key = key2

            orig_con_dict[float(key)] = [float(min_key), min_diff]

    # Getting the average first 2 eigenvector components for each of the regions
    for key, region in eig_dict.iteritems():
        #         print(key)
        y_vals = []

        for j in range(len(region)):
            y_vals.append(region[j])
        y_vals = np.array(y_vals)
        x_avg = np.mean(y_vals[:, 0])
        y_avg = np.mean(y_vals[:, 1])
        z_avg = np.mean(y_vals[:, 2])
        avg_dict[key] = [x_avg, y_avg, z_avg]

    # print('getcon avg_dict')
    #     print(avg_dict)

    # Computing connectivity between regions using the distance between averages
    for key, avg in avg_dict.iteritems():
        min_key = ''
        min_diff = float('inf')
        for key2, avg2 in avg_dict.iteritems():
            if key2 == key:
                continue
            avg_np = np.array(avg)
            avg2_np = np.array(avg2)
            diff = np.linalg.norm(avg_np - avg2_np)
            if diff < min_diff:
                min_diff = diff
                min_key = key2

        con_dict[float(key)] = [float(min_key), min_diff]

    con_dict = OrderedDict(sorted(con_dict.items()))
    orig_con_dict = OrderedDict(sorted(orig_con_dict.items()))

    g = create_connectivity_graph(orig_avg_dict, avg_dict, max_dist)

    if orig_dict == None:
        return con_dict
    else:
        return con_dict, orig_con_dict, g

class SparseMatrix:
    def __init__(self, x, y, z):
#         self._max_index = 0
        x_dim = x
        y_dim = y
        z_dim = z
        self._vector = {}

    def add(self, index, value):
        # vector starts at index one, because it reads from the file and the file
        # always has the index of the features start at 1
        self._vector[index] = value
#         if index > self._max_index:
#             self._max_index = index

    def get(self, index):
        # if the index doesn't exist in the dict, return 0 because it's sparse anyways
        if index in self._vector:
            return self._vector[index]
        return -1

    def get_sparse_matrix(self):
        return self._vector
        # return self._vector.keys()

#     def get_full_vector(self, size=None):
#         """ Returns a full vector of features as a numpy array. """
#         size = (self._max_index + 1) if size == None else size
#         full_vector = np.zeros(size)  # 0 indexed
#         for key, value in self._vector.iteritems():
#             full_vector[key] = value

#         return full_vector

    def __str__(self):
        return str(self._vector)

def plot_con_mat(con_adj_mat, output_path=None, show=False):
    title = 'Connectivity Heatmap'
    data = [
        Heatmap(
            z = con_adj_mat,
    #         x = con_graph.nodes(),
    #         y = con_graph.nodes()
        )
    ]
    layout = Layout(
        title = title,
        xaxis=dict(title='region'),
        yaxis=dict(title='region')
    )
    fig = Figure(data=data, layout=layout)
    if show:
        iplot(fig)
    if output_path != None:
        plotly.offline.plot(fig, filename=output_path)

    return fig


# Author: Vlad Niculae
# Licence: BSD 3 clause

import numpy as np

from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_warns
from sklearn.utils.testing import ignore_warnings

from sklearn.linear_model import (orthogonal_mp, orthogonal_mp_gram,
                                  OrthogonalMatchingPursuit,
                                  OrthogonalMatchingPursuitCV,
                                  LinearRegression)
from sklearn.utils import check_random_state
from sklearn.datasets import make_sparse_coded_signal

n_samples, n_features, n_nonzero_coefs, n_targets = 20, 30, 5, 3
y, X, gamma = make_sparse_coded_signal(n_targets, n_features, n_samples,
                                       n_nonzero_coefs, random_state=0)
G, Xy = np.dot(X.T, X), np.dot(X.T, y)
# this makes X (n_samples, n_features)
# and y (n_samples, 3)

def test_correct_shapes():
    assert_equal(orthogonal_mp(X, y[:, 0], n_nonzero_coefs=5).shape,
                 (n_features,))
    assert_equal(orthogonal_mp(X, y, n_nonzero_coefs=5).shape,
                 (n_features, 3))

def test_correct_shapes_gram():
    assert_equal(orthogonal_mp_gram(G, Xy[:, 0], n_nonzero_coefs=5).shape,
                 (n_features,))
    assert_equal(orthogonal_mp_gram(G, Xy, n_nonzero_coefs=5).shape,
                 (n_features, 3))

def test_n_nonzero_coefs():
    assert_true(np.count_nonzero(orthogonal_mp(X, y[:, 0],
                                 n_nonzero_coefs=5)) <= 5)
    assert_true(np.count_nonzero(orthogonal_mp(X, y[:, 0], n_nonzero_coefs=5,
                                               precompute=True)) <= 5)

def test_tol():
    tol = 0.5
    gamma = orthogonal_mp(X, y[:, 0], tol=tol)
    gamma_gram = orthogonal_mp(X, y[:, 0], tol=tol, precompute=True)
    assert_true(np.sum((y[:, 0] - np.dot(X, gamma)) ** 2) <= tol)
    assert_true(np.sum((y[:, 0] - np.dot(X, gamma_gram)) ** 2) <= tol)

def test_with_without_gram():
    assert_array_almost_equal(
        orthogonal_mp(X, y, n_nonzero_coefs=5),
        orthogonal_mp(X, y, n_nonzero_coefs=5, precompute=True))

def test_with_without_gram_tol():
    assert_array_almost_equal(
        orthogonal_mp(X, y, tol=1.),
        orthogonal_mp(X, y, tol=1., precompute=True))

def test_unreachable_accuracy():
    assert_array_almost_equal(
        orthogonal_mp(X, y, tol=0),
        orthogonal_mp(X, y, n_nonzero_coefs=n_features))

    assert_array_almost_equal(
        assert_warns(RuntimeWarning, orthogonal_mp, X, y, tol=0,
                     precompute=True),
        orthogonal_mp(X, y, precompute=True,
                      n_nonzero_coefs=n_features))

def test_bad_input():
    assert_raises(ValueError, orthogonal_mp, X, y, tol=-1)
    assert_raises(ValueError, orthogonal_mp, X, y, n_nonzero_coefs=-1)
    assert_raises(ValueError, orthogonal_mp, X, y,
                  n_nonzero_coefs=n_features + 1)
    assert_raises(ValueError, orthogonal_mp_gram, G, Xy, tol=-1)
    assert_raises(ValueError, orthogonal_mp_gram, G, Xy, n_nonzero_coefs=-1)
    assert_raises(ValueError, orthogonal_mp_gram, G, Xy,
                  n_nonzero_coefs=n_features + 1)

def test_perfect_signal_recovery():
    idx, = gamma[:, 0].nonzero()
    gamma_rec = orthogonal_mp(X, y[:, 0], 5)
    gamma_gram = orthogonal_mp_gram(G, Xy[:, 0], 5)
    assert_array_equal(idx, np.flatnonzero(gamma_rec))
    assert_array_equal(idx, np.flatnonzero(gamma_gram))
    assert_array_almost_equal(gamma[:, 0], gamma_rec, decimal=2)
    assert_array_almost_equal(gamma[:, 0], gamma_gram, decimal=2)

def test_estimator():
    omp = OrthogonalMatchingPursuit(n_nonzero_coefs=n_nonzero_coefs)
    omp.fit(X, y[:, 0])
    assert_equal(omp.coef_.shape, (n_features,))
    assert_equal(omp.intercept_.shape, ())
    assert_true(np.count_nonzero(omp.coef_) <= n_nonzero_coefs)

    omp.fit(X, y)
    assert_equal(omp.coef_.shape, (n_targets, n_features))
    assert_equal(omp.intercept_.shape, (n_targets,))
    assert_true(np.count_nonzero(omp.coef_) <= n_targets * n_nonzero_coefs)

    omp.set_params(fit_intercept=False, normalize=False)

    omp.fit(X, y[:, 0])
    assert_equal(omp.coef_.shape, (n_features,))
    assert_equal(omp.intercept_, 0)
    assert_true(np.count_nonzero(omp.coef_) <= n_nonzero_coefs)

    omp.fit(X, y)
    assert_equal(omp.coef_.shape, (n_targets, n_features))
    assert_equal(omp.intercept_, 0)
    assert_true(np.count_nonzero(omp.coef_) <= n_targets * n_nonzero_coefs)

def test_identical_regressors():
    newX = X.copy()
    newX[:, 1] = newX[:, 0]
    gamma = np.zeros(n_features)
    gamma[0] = gamma[1] = 1.
    newy = np.dot(newX, gamma)
    assert_warns(RuntimeWarning, orthogonal_mp, newX, newy, 2)

def test_swapped_regressors():
    gamma = np.zeros(n_features)
    # X[:, 21] should be selected first, then X[:, 0] selected second,
    # which will take X[:, 21]'s place in case the algorithm does
    # column swapping for optimization (which is the case at the moment)
    gamma[21] = 1.0
    gamma[0] = 0.5
    new_y = np.dot(X, gamma)
    new_Xy = np.dot(X.T, new_y)
    gamma_hat = orthogonal_mp(X, new_y, 2)
    gamma_hat_gram = orthogonal_mp_gram(G, new_Xy, 2)
    assert_array_equal(np.flatnonzero(gamma_hat), [0, 21])
    assert_array_equal(np.flatnonzero(gamma_hat_gram), [0, 21])

def test_no_atoms():
    y_empty = np.zeros_like(y)
    Xy_empty = np.dot(X.T, y_empty)
    gamma_empty = ignore_warnings(orthogonal_mp)(X, y_empty, 1)
    gamma_empty_gram = ignore_warnings(orthogonal_mp)(G, Xy_empty, 1)
    assert_equal(np.all(gamma_empty == 0), True)
    assert_equal(np.all(gamma_empty_gram == 0), True)

def test_omp_path():
    path = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=True)
    last = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=False)
    assert_equal(path.shape, (n_features, n_targets, 5))
    assert_array_almost_equal(path[:, :, -1], last)
    path = orthogonal_mp_gram(G, Xy, n_nonzero_coefs=5, return_path=True)
    last = orthogonal_mp_gram(G, Xy, n_nonzero_coefs=5, return_path=False)
    assert_equal(path.shape, (n_features, n_targets, 5))
    assert_array_almost_equal(path[:, :, -1], last)

def test_omp_return_path_prop_with_gram():
    path = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=True,
                         precompute=True)
    last = orthogonal_mp(X, y, n_nonzero_coefs=5, return_path=False,
                         precompute=True)
    assert_equal(path.shape, (n_features, n_targets, 5))
    assert_array_almost_equal(path[:, :, -1], last)

def test_omp_cv():
    y_ = y[:, 0]
    gamma_ = gamma[:, 0]
    ompcv = OrthogonalMatchingPursuitCV(normalize=True, fit_intercept=False,
                                        max_iter=10, cv=5)
    ompcv.fit(X, y_)
    assert_equal(ompcv.n_nonzero_coefs_, n_nonzero_coefs)
    assert_array_almost_equal(ompcv.coef_, gamma_)
    omp = OrthogonalMatchingPursuit(normalize=True, fit_intercept=False,
                                    n_nonzero_coefs=ompcv.n_nonzero_coefs_)
    omp.fit(X, y_)
    assert_array_almost_equal(ompcv.coef_, omp.coef_)

def test_omp_reaches_least_squares():
    # Use small simple data; it's a sanity check but OMP can stop early
    rng = check_random_state(0)
    n_samples, n_features = (10, 8)
    n_targets = 3
    X = rng.randn(n_samples, n_features)
    Y = rng.randn(n_samples, n_targets)
    omp = OrthogonalMatchingPursuit(n_nonzero_coefs=n_features)
    lstsq = LinearRegression()
    omp.fit(X, Y)
    lstsq.fit(X, Y)
    assert_array_almost_equal(omp.coef_, lstsq.coef_)