#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2020 Félix Chénier

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Provide functions to process kinetic data from instrumented wheelchair wheels.
"""

__author__ = "Félix Chénier"
__copyright__ = "Copyright (C) 2020 Félix Chénier"
__email__ = "chenier.felix@uqam.ca"
__license__ = "Apache 2.0"

import kineticstoolkit.filters as filters
import kineticstoolkit.cycles as cycles
from kineticstoolkit import TimeSeries
from kineticstoolkit.decorators import unstable, dead, directory

import numpy as np
from numpy import sin, cos, pi
import pandas as pd
import warnings
import struct  # to unpack binary data from SmartWheels' txt files
from typing import Union, Optional, List

def read_file(filename: str, /, file_format: str = '') -> TimeSeries:
    """
    Read a file containing pushrim kinetics data.

    Parameters
    ----------
    filename
        Name of the file to open
    file_format
        Format of the file. Can be either:

        - 'smartwheel' (SmartWheel CSV file)
        - 'smartwheeltxt' (SmartWheel SD-Card TXT file)
        - 'racingwheel' (will change)

    """
    if file_format == '':
        warnings.warn("file_format will need to be explicitely specified in "
                      "future versions. Now using 'smartwheel'.",
                      FutureWarning)
        file_format = 'smartwheel'

    if file_format == 'smartwheel':

        dataframe = pd.read_csv(filename, delimiter=None, header=None)
        if dataframe.shape[1] == 1:  # Retry with ; as separator
            dataframe = pd.read_csv(filename, delimiter=';', header=None)

        data = dataframe.to_numpy()
        index = data[:, 1]
        time = np.arange(0, len(index)) / 240
        channels = data[:, 6:12]
        forces = data[:, 18:21]
        moments = data[:, 21:24]
        angle_deg = data[:, 3]
        angle_rad = np.unwrap(np.deg2rad(angle_deg))

        ts = TimeSeries(time=time)

        ts.data['Index'] = index
        ts.data['Channels'] = channels
        ts.data['Forces'] = np.block([[forces, np.zeros((len(index), 1))]])
        ts.data['Moments'] = np.block([[moments, np.zeros((len(index), 1))]])
        ts.data['Angle'] = angle_rad

        ts.add_data_info('Channels', 'Unit', 'raw')
        ts.add_data_info('Forces', 'Unit', 'N')
        ts.add_data_info('Moments', 'Unit', 'Nm')
        ts.add_data_info('Angle', 'Unit', 'rad')

    elif file_format == 'racingwheel':

        dataframe = pd.read_csv(filename, delimiter=',')
        data = dataframe.to_numpy()
        time = data[:, 0]
        channels = data[:, 1:7]
        battery = data[:, 8]

        ts = TimeSeries(time=time)

        ts.data['Channels'] = channels
        ts.add_data_info('Channels', 'Unit', 'raw')

        ts.data['Battery'] = battery
        ts.add_data_info('Battery', 'Unit', 'raw')

    elif file_format == 'smartwheeltxt':

        data = {'ch1': [], 'ch2': [], 'ch3': [],
                'ch4': [], 'ch5': [], 'ch6': [], 'angle_ticks': []}

        length = 0
        with open(filename, 'rb') as fid:

            while True:
                try:
                    _ = fid.read(2)
                    data['ch1'].append(struct.unpack('h', fid.read(2))[0])
                    data['ch2'].append(struct.unpack('h', fid.read(2))[0])
                    data['ch3'].append(struct.unpack('h', fid.read(2))[0])
                    data['ch4'].append(struct.unpack('h', fid.read(2))[0])
                    data['ch5'].append(struct.unpack('h', fid.read(2))[0])
                    data['ch6'].append(struct.unpack('h', fid.read(2))[0])
                    data['angle_ticks'].append(
                        struct.unpack('i', fid.read(4))[0])
                    _ = fid.read(8)
                    length += 1
                except Exception:
                    break

        ch1 = np.array(data['ch1'][1:length])  # Remove 1st sample to be
        ch2 = np.array(data['ch2'][1:length])  # consistent with CSV file
        ch3 = np.array(data['ch3'][1:length])
        ch4 = np.array(data['ch4'][1:length])
        ch5 = np.array(data['ch5'][1:length])
        ch6 = np.array(data['ch6'][1:length])
        angle_ticks = np.array(data['angle_ticks'][1:length])

        # Keep only 12 least significant bytes
        ch1 = np.mod(ch1, 2 ** 12)
        ch2 = np.mod(ch2, 2 ** 12)
        ch3 = np.mod(ch3, 2 ** 12)
        ch4 = np.mod(ch4, 2 ** 12)
        ch5 = np.mod(ch5, 2 ** 12)
        ch6 = np.mod(ch6, 2 ** 12)

        # Convert angle in radian
        angle = angle_ticks / 4096 * 2 * np.pi

        ts = TimeSeries(
            time=np.linspace(0, (length - 1) / 240, length - 1))
        ts.data['Channels'] = np.concatenate(
            [ch1[:, np.newaxis],
             ch2[:, np.newaxis],
             ch3[:, np.newaxis],
             ch4[:, np.newaxis],
             ch5[:, np.newaxis],
             ch6[:, np.newaxis]], axis=1)
        ts.data['Angle'] = angle

        ts.add_data_info('Channels', 'Unit', 'raw')
        ts.add_data_info('Angle', 'Unit', 'rad')

    else:
        raise ValueError('Unknown file format.')

    return ts

@unstable
def find_recovery_indices(Mz: np.ndarray, /) -> np.ndarray:
    """
    Find recovery indices based on a vector of propulsion moments.

    This function analyzes the Mz moments to find which data correspond to
    pushes and which data correspond to recoveries. The method is very
    conservative on what could be considered as a recovery, so that every
    index returned by this function is almost certain to correspond to a
    recovery. This function is used by `pushrimkinetics.remove_sinusoids`
    to identify the instants with no hand contact. It should not be used to
    isolate the push and recovery phases (use `ktk.cycles.detect_cycles()`
    instead).

    Parameters
    ----------
    Mz
        Array that contains the propulsion moments in Nm.

    Returns
    -------
    np.ndarray
        Array of bools where each True represents recovery.

    See Also
    --------
    ktk.cycles.detect_cycles

    """
    Mz = Mz.copy()

    threshold = 2.24  # (Nm): max tolerance for the remaining values.

    while np.nanmax(Mz) - np.nanmin(Mz) > threshold:

        # Remove 1% of data that are the farthest to the median:

        # Sort data
        index_to_remove = np.argsort(np.abs(Mz - np.nanmedian(Mz)))
        sorted_Mz = Mz[index_to_remove]
        index_to_remove = index_to_remove[~np.isnan(sorted_Mz)]

        # Remove the 1% upper.
        index_to_remove = index_to_remove[
            int(0.99*len(index_to_remove))-1:]

        # Assign nan to these data
        Mz[index_to_remove] = np.nan

    index = ~np.isnan(Mz)

    return index

def remove_offsets(
        kinetics: TimeSeries,
        baseline_kinetics: Optional[TimeSeries] = None
) -> TimeSeries:
    """
    Remove dynamic offsets in forces and moments.

    Parameters
    ----------
    kinetics
        TimeSeries that contains at least Forces, Moments and Angle data.
    baseline_kinetics
        Optional. TimeSeries that contains at least Forces and Moments data.
        This TimeSeries contains a baseline trial, where the wheelchair must be
        pushed by an operator and where no external force must be applied on
        the pushrims. If no baseline is provided, the baseline is calculated
        based on a detection of recoveries in the supplied kinetics
        TimeSeries.

    Returns
    -------
    TimeSeries
        A copy of the input TimeSeries, where sinusoids are removed from
        Forces and Moments data.

    References
    ----------
    F. Chénier, R. Aissaoui, C. Gauthier, and D. H. Gagnon,
    "Wheelchair pushrim kinetics measurement: A method to cancel
    inaccuracies due to pushrim weight and wheel camber," Medical
    Engineering and Physics, vol. 40, pp. 75--86, 2017.

    """
    kinetics = kinetics.copy()

    if baseline_kinetics is None:
        # Create baseline kinetics.
        recovery_index = find_recovery_indices(kinetics.data['Moments'][:, 2])
        f_ofs = np.hstack((kinetics.data['Forces'][recovery_index, 0:3],
                           kinetics.data['Moments'][recovery_index, 0:3]))
        theta_baseline = kinetics.data['Angle'][recovery_index]

    else:
        # Use baseline kinetics.
        f_ofs = np.hstack((baseline_kinetics.data['Forces'][:, 0:3],
                           baseline_kinetics.data['Moments'][:, 0:3]))
        theta_baseline = baseline_kinetics.data['Angle'][:]

    # Do the regression
    theta_baseline = theta_baseline[:, np.newaxis]
    q = np.hstack((
        np.sin(theta_baseline),
        np.cos(theta_baseline),
        np.ones((len(theta_baseline), 1))
    ))
    A = np.linalg.lstsq(q, f_ofs, rcond=None)
    A = A[0]

    # Apply the regression to forces and moments
    theta = kinetics.data['Angle']
    theta = theta[:, np.newaxis]

    f = np.hstack((kinetics.data['Forces'][:, 0:3],
                   kinetics.data['Moments'][:, 0:3]))

    q = np.hstack((
        np.sin(theta),
        np.cos(theta),
        np.ones((len(theta), 1))
    ))

    f = f - q @ A

    # Make the output timeseries
    kinetics.data['Forces'][:, 0:3] = f[:, 0:3]
    kinetics.data['Moments'][:, 0:3] = f[:, 3:6]

    return kinetics

def calculate_forces_and_moments(
        kinetics: TimeSeries, /,
        gains: Union[np.ndarray, str],
        offsets: np.ndarray = np.zeros((6)), *,
        transducer: str = 'force_cell',
        reference_frame: str = 'wheel') -> TimeSeries:
    """
    Calculate pushrim forces and moments based on raw channel values.

    For standard force cells (with each channel being a raw value
    corresponding to Fx, Fy, Fz, Mx, My, Mz, respectively), calculates
    the forces and moments using a sensitivity matrix (gains) and an
    offset vector (offsets):

    ``[Fx, Fy, Fz, Mx, My, Mz] = gains @ channels + offsets``

    For SmartWheel, calculates the forces and moments using a gain
    vector (gains) and an offset vector (offsets).

    Parameters
    ----------
    kinetics
        Input TimeSeries that must contain a 'Channels' key in its data dict.
    gains
        6x6 gain matrix (force_cell) or gain vector of length 6 (smartwheel).
    offsets
        Optional. Offset vector of length 6.
    transducer
        Optional. 'force_cell' or 'smartwheel'.
    reference_frame
        Optional. 'wheel' or 'hub'. 'wheel' to report the forces and moments
        into the local wheel's reference frame; 'hub' to compensate for the
        wheel rotation and match the reference frame used by the SmartWheel:
        x anteroposterior, y in the wheel plane, upward for non-camberred
        wheels, and z perpendicular to the wheel plane, outward.

    Returns
    -------
    TimeSeries
        A copy of the input TimeSeries, with the added 'Forces'
        and 'Moments' data keys.

    Note
    ----
    Some calibration matrices are provided as examples in the
    ``pushrimkinetics.CALIBRATION_MATRICES`` dictionary. This dictionary can
    be used directly using dict unpacking. For example::

        ktk.pushrimkinetics.calculate_force_and_moments(
            kinetics,
            **ktk.pushrimkinetics.CALIBRATION_MATRICES['SmartWheel_123'],
            reference_frame='level'
        )

    """
    # Calculate the forces and moments and add to the output
    if transducer == 'smartwheel':

        # Calculate the rotation angle to apply to the calculated kinetics
        if reference_frame == 'wheel':
            theta = np.array([0.0])
        elif reference_frame == 'hub':
            theta = kinetics.data['Angle']
        else:
            raise ValueError("reference_frame must be 'wheel' or 'hub'")

        # Extract channels and angle
        ch = kinetics.data['Channels'] - 2048

        # Calculate the forces and moments
        Fx = gains[0] * (
            ch[:, 0] * sin(theta) +
            ch[:, 2] * sin(theta + 2 * pi / 3) +
            ch[:, 4] * sin(theta + 4 * pi / 3)) + offsets[0]

        Fy = gains[1] * (
            ch[:, 0] * cos(theta) +
            ch[:, 2] * cos(theta + 2 * pi / 3) +
            ch[:, 4] * cos(theta + 4 * pi / 3)) + offsets[1]

        Fz = gains[2] * (ch[:, 1] + ch[:, 3] + ch[:, 5]) + offsets[2]

        Mx = gains[3] * (
            ch[:, 1] * sin(theta) +
            ch[:, 3] * sin(theta + 2 * pi / 3) +
            ch[:, 5] * sin(theta + 4 * pi / 3)) + offsets[3]

        My = gains[4] * (
            ch[:, 1] * cos(theta) +
            ch[:, 3] * cos(theta + 2 * pi / 3) +
            ch[:, 5] * cos(theta + 4 * pi / 3)) + offsets[4]

        Mz = gains[5] * (ch[:, 0] + ch[:, 2] + ch[:, 4]) + offsets[5]
        forces_moments = np.block([Fx[:, np.newaxis],
                                   Fy[:, np.newaxis],
                                   Fz[:, np.newaxis],
                                   Mx[:, np.newaxis],
                                   My[:, np.newaxis],
                                   Mz[:, np.newaxis]])

    elif transducer == 'force_cell':

        # Calculate the rotation angle to apply to the calculated kinetics
        if reference_frame == 'wheel':
            theta = np.array([0.0])
        elif reference_frame == 'hub':
            raise NotImplementedError("hub reference_frame not implemented yet"
                                      "for force_cell transducers.")
        else:
            raise ValueError("reference_frame must be 'wheel' or 'hub'")

        n_frames = kinetics.data['Channels'].shape[0]

        forces_moments = np.empty((n_frames, 6))
        for i_frame in range(n_frames):
            forces_moments[i_frame] = (gains @
                                       kinetics.data['Channels'][i_frame] +
                                       offsets)

    # Format these data in the output timeseries
    kinetics = kinetics.copy()

    kinetics.data['Forces'] = np.concatenate(
        [forces_moments[:, 0:3], np.zeros((forces_moments.shape[0], 1))],
        axis=1)
    kinetics.add_data_info('Forces', 'Unit', 'N')

    kinetics.data['Moments'] = np.concatenate(
        [forces_moments[:, 3:6], np.zeros((forces_moments.shape[0], 1))],
        axis=1)
    kinetics.add_data_info('Moments', 'Unit', 'Nm')

    return(kinetics)

def calculate_velocity(tsin: TimeSeries, /) -> TimeSeries:
    """
    Calculate velocity based on wheel angle.

    The velocity is calculated by deriving the angle using a 2nd order
    Savitzky-Golay filter of length 21. This filter has been experimentally
    validated to maximize the signal-to-noise ratio for a SmartWheel recording
    at 240 Hz. This function may change signature in the future to include
    other filtering options. To manually get the velocity using a custom
    filter, please see the ``ktk.filters`` module.

    Parameters
    ----------
    tsin
        TimeSeries that contains at least the data key 'Angle'.

    Returns
    -------
    TimeSeries
                A copy of the TimeSeries with the added data key 'Velocity'.

    See Also
    --------
    ktk.filters.butter : Butterworth filter for TimeSeries
    ktk.filters.savgol : Savitsky-golay filter for TimeSeries
    ktk.filters.deriv : Derivative filter for TimeSeries

    """
    tsangle = TimeSeries()
    tsangle.time = tsin.time
    tsangle.data['Angle'] = tsin.data['Angle']
    tsvelocity = filters.savgol(tsangle, window_length=21,
                                poly_order=2, deriv=1)
    tsout = tsin.copy()
    tsout.data['Velocity'] = tsvelocity.data['Angle']
    try:
        tsout.add_data_info('Velocity', 'Unit',
                            tsout.data_info['Angle']['Unit'] + '/s')
    except KeyError:
        pass
    return tsout

def calculate_power(tsin: TimeSeries, /) -> TimeSeries:
    """
    Calculate power based on wheel velocity and moment.

    Parameters
    ----------
    tsin
        TimeSeries that contains at least the data keys 'Velocity' and
        'Moments'. The units must be consistent (e.g., rad/s and Nm)

    Returns
    -------
    TimeSeries
        A copy of the TimeSeries with the added data key 'Power'.

    """
    tsout = tsin.copy()
    tsout.data['Power'] = (tsout.data['Velocity'] *
                           tsout.data['Moments'][:, 2])
    tsout.add_data_info('Power', 'Unit', 'W')
    return tsout

#--- Deprecated functions ---#
@dead(since="October 2020",
      until="December 2021",
      details="Please use ktk.cycles.detect_cycles instead.")
def detect_pushes(
        tsin: TimeSeries, /, *,
        push_threshold: float = 5.0,
        recovery_threshold: float = 2.0,
        min_push_time: float = 0.1,
        min_push_force: float = 30.0) -> TimeSeries:
    """
    Detect pushes and recoveries automatically.

    Parameters
    ----------
    tsin
        Input TimeSeries that must contain a 'Forces' key in its data dict.
    push_threshold
        Optional. The total force over which a push phase is triggered, in
        newton.
    recovery_threshold
        Optional. The total force under which a recovery phase is triggered,
        in newton.
    min_push_time
        Optional. The minimum time required for a push time, in seconds.
        Detected pushes that last less than this minimum time are removed from
        the push analysis.
    min_recovery_time
        Optional. The minimum time required for a recovery time, in seconds.
        Detected recoveries that last less than this minimum time are removed
        from the push analysis.
    min_push_force
        Optional. The minimum total push force in N under which the detected
        push is discarded. For example, if the user puts their hands on the
        pushrim before starting propelling, this may be detected as a push.
        Using a minimum push force removes these misdetected pushes.

    Returns
    -------
    TimeSeries
        A copy of tsin with the following added events:
        - 'push'
        - 'recovery'

    """
    # Calculate the total force
    f_tot = np.sqrt(np.sum(tsin.data['Forces']**2, axis=1))
    ts_force = TimeSeries(time=tsin.time, data={'Ftot': f_tot})
    ts_force.events = tsin.events

    # Smooth the total force to avoid detecting pushes on glitches
    ts_force = filters.smooth(ts_force, 11)

    # Remove the median if it existed
    ts_force.data['Ftot'] = \
        ts_force.data['Ftot'] - np.median(ts_force.data['Ftot'])

    # Find the pushes
    ts_force = cycles.detect_cycles(
        ts_force, 'Ftot',
        event_name1='push',
        event_name2='recovery',
        threshold1=push_threshold,
        threshold2=recovery_threshold,
        min_duration1=min_push_time,
        min_peak_height1=min_push_force)

    # Form the output timeseries
    tsout = tsin.copy()
    tsout.events = ts_force.events

    return tsout

@dead(since="October 2020",
      until="December 2021",
      details="Please use ktk.pushrimkinetics.remove_offsets instead.")
def remove_sinusoids(
        kinetics: TimeSeries,
        baseline_kinetics: Optional[TimeSeries] = None) -> TimeSeries:
    """Remove offsets."""
    return remove_offsets(kinetics, baseline_kinetics)

#--- Some calibration matrices ---#
CALIBRATION_MATRICES = {}
CALIBRATION_MATRICES['SmartWheel_93'] = {
    'gains': np.array([-0.1080, 0.1080, 0.0930, 0.0222, -0.0222, 0.0234999]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_94'] = {
    'gains': np.array([-0.1070, 0.1070, 0.0960, 0.0222, -0.0222, 0.0230]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_123'] = {
    'gains': np.array([-0.106, 0.106, 0.094, 0.022, -0.022, 0.0234999]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_124'] = {
    'gains': np.array([-0.106, 0.106, 0.0949999, 0.0215, -0.0215, 0.0225]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_125'] = {
    'gains': np.array([-0.104, 0.104, 0.0979999, 0.0215, -0.0215, 0.0225]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_126'] = {
    'gains': np.array([-0.1059999, 0.1059999, 0.086, 0.021, -0.021, 0.023]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_126_S18'] = {
    'gains': np.array([-0.1083, 0.1109, 0.0898, 0.0211, -0.0194, 0.0214]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_179_S18'] = {
    'gains': np.array([-0.1399, 0.1091, 0.0892, 0.0240, -0.0222, 0.0241]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_180_S18'] = {
    'gains': np.array([-0.1069, 0.1091, 0.0932, 0.0240, -0.0226, 0.0238]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['SmartWheel_181_S18'] = {
    'gains': np.array([-0.1152, 0.1095, 0.0791, 0.0229, -0.0197, 0.0220]),
    'offsets': np.array([0.0, 10.0, 0.0, 0.0, 0.0, 0.0]),
    'transducer': 'smartwheel',
}
CALIBRATION_MATRICES['MOSA_Racing_1'] = {
    'gains': (np.array([
        [201.027, 1.387, 2.077, -3.852, -1.837, -1.519],
        [-0.840, 201.396, 2.119, 0.083, -6.877, 4.482],
        [-1.935, -1.643, 402.286, 1.687, 0.897, -23.616],
        [0.213, 0.122, 0.120, 25.190, -0.013, 0.147],
        [-0.072, 0.286, 0.076, 0.012, 25.430, 0.146],
        [0.016, -0.015, 0.046, -0.099, -0.076, 25.206]])  # Cell calibration
        / (2.**15) / 10  # ADC gains
        / np.array([-2., -2., -2., -2., -4., -4.])  # Board gains
    ),
    'offsets': [-111.3874, -63.3298, -8.6596, 1.8089, 1.5761, -0.8869],
    'transducer': 'force_cell',
}

module_locals = locals()

def __dir__():  # pragma: no cover
    return directory(module_locals)

if __name__ == "__main__":  # pragma: no cover
    import doctest
    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)


"""
Module for rover perception.

Contains functions for processing rover's front camera image frames
and updating rover state.

Objective:
Turn rover camera 3D images into a 2D perspective world-view of the
rover environment that identifies regions of interests (ROIs), and
superimpose this view on the ground truth worldmap.

NOTE:

Units:
time -- seconds
distance -- meters
velocity -- meters/second
angle, heading -- degrees
yaw, pitch, roll -- degrees

Short Forms:
pixpts -- pixel points
nav -- navigable terrain pixels
obs -- obstacle pixels
rock -- rock pixels

Abbreviations:
ROI -- Regions of interest
pf -- perspective frame
rf -- rover frame
wf -- world frame

"""

__author__ = 'Salman Hashmi, Ryan Keenan'
__license__ = 'BSD License'

from collections import namedtuple

import numpy as np
import cv2

def color_thresh(input_img, rgb_thresh=(160, 160, 160),
                 low_bound=(75, 130, 130), upp_bound=(255, 255, 255)):
    """
    Apply color thresholds to extract pixels of navigable/obstacles/rocks.

    Keyword arguments:
    input_img -- numpy image on which RGB threshold is applied
    rgb_thresh -- RGB thresh tuple above which only ground pixels are detected
    low/up_bounds -- HSV tuples defining color range of gold rock samples

    Return value:
    thresh_imgs -- namedtuple of binary images identifying nav/obs/rock pixels

    """
    # Create arrays of zeros same xy size as input_img, but single channel
    nav_img = np.zeros_like(input_img[:, :, 0])
    obs_img = np.zeros_like(input_img[:, :, 0])

    # Convert BGR input_img to HSV for rock samples
    hsv_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2HSV)

    # Require that each of the R(0), G(1), B(2) pixels be above all three
    # rgb_thresh values such that pixpts_above_thresh will now contain a
    # boolean array with "True" where threshold was met
    pixpts_above_thresh = ((input_img[:, :, 0] > rgb_thresh[0])
                           & (input_img[:, :, 1] > rgb_thresh[1])
                           & (input_img[:, :, 2] > rgb_thresh[2]))

    pixpts_nonzero = ((input_img[:, :, 0] > 0)
                      & (input_img[:, :, 1] > 0)
                      & (input_img[:, :, 2] > 0))

    # obstacle pixels are those non-zero pixels where rgb_thresh was not met
    obs_pixpts = np.logical_and(
        pixpts_nonzero, np.logical_not(pixpts_above_thresh)
        )
    # Index the array of zeros with the boolean array and set to 1
    # those pixels where ROI threshold was met
    nav_img[pixpts_above_thresh] = 1
    obs_img[obs_pixpts] = 1

    # Threshold the HSV image to get only colors for gold rock samples
    rock_img = cv2.inRange(hsv_img, low_bound, upp_bound)

    # Return the threshed binary images
    ThreshedImages = namedtuple('ThreshedImages', 'nav obs rock')
    thresh_imgs = ThreshedImages(nav_img, obs_img, rock_img)

    return thresh_imgs

def perspect_transform(src_img, dst_grid=10, bottom_offset=6):
    """
    Apply a perspective transformation to input 3D image.

    Keyword arguments:
    src_img -- 3D numpy image on which perspective transform is applied
    dst_grid -- size of 2D output image box of 10x10 pixels equaling 1 Sq m
    bottom_offset -- bottom of cam image is some distance in front of rover

    Return value:
    dst_img -- 2D warped numpy image with overhead view

    """
    # Dimension of source image from rover camera
    height, width = src_img.shape[0], src_img.shape[1]

    # Numpy array of four source points defining a grid on input 3D image
    # acquired from calibration data in test notebook
    src_x1, src_y1 = 14, 140
    src_x2, src_y2 = 301, 140
    src_x3, src_y3 = 200, 96
    src_x4, src_y4 = 118, 96

    # Corresponding destination points on output 2D overhead image
    dst_x1, dst_y1 = (width/2 - dst_grid/2), (height-bottom_offset)
    dst_x2, dst_y2 = (width/2 + dst_grid/2), (height-bottom_offset)
    dst_x3, dst_y3 = (width/2 + dst_grid/2), (height-dst_grid-bottom_offset)
    dst_x4, dst_y4 = (width/2 - dst_grid/2), (height-dst_grid-bottom_offset)

    src_points_3d = np.float32([[src_x1, src_y1],
                                [src_x2, src_y2],
                                [src_x3, src_y3],
                                [src_x4, src_y4]])

    dst_points_2d = np.float32([[dst_x1, dst_y1],
                                [dst_x2, dst_y2],
                                [dst_x3, dst_y3],
                                [dst_x4, dst_y4]])

    transform_matrix = cv2.getPerspectiveTransform(src_points_3d,
                                                   dst_points_2d)
    # Keep same size as source image
    dst_img = cv2.warpPerspective(src_img, transform_matrix, (width, height))

    return dst_img

def perspect_to_rover(binary_img):
    """
    Transform pixel points from perspective frame to rover frame.

    Keyword arguments:
    binary_img -- single channel 2D warped numpy image in perspective frame

    Return value:
    pixpts_rf -- tuple of numpy arrays of pixel x,y points in rover frame

    """
    # Dimension of input image
    height, width = binary_img.shape[0], binary_img.shape[1]

    # Identify all nonzero pixel coords in the binary image
    ypix_pts_pf, xpix_pts_pf = binary_img.nonzero()

    # Calculate pixel positions with reference to rover's coordinate
    # frame given that rover front camera itself is at center bottom
    # of the photographed image
    xpix_pts_rf = -(ypix_pts_pf - height).astype(np.float)
    ypix_pts_rf = -(xpix_pts_pf - width/2).astype(np.float)
    pixpts_rf = xpix_pts_rf, ypix_pts_rf

    return pixpts_rf

def to_polar_coords(pixpts):
    """
    Convert cartesian coordinates of pixels to polar coordinates.

    Keyword arguments:
    pixpts -- tuple of numpy arrays of pixel x,y points

    Return value:
    dists, angles -- distance(m) and angles(deg) to pixpts

    """
    rad2deg = 180./np.pi
    xpix_pts, ypix_pts = pixpts

    dists = np.sqrt(xpix_pts**2 + ypix_pts**2)
    angles = np.arctan2(ypix_pts, xpix_pts)*rad2deg

    return dists, angles

def rotate_pixpts(pixpts, angle):
    """
    Geometrically rotate pixel points by specified angle.

    Keyword arguments:
    pixpts -- tuple of numpy arrays of pixel x,y points
    angle -- rotation angle

    Return value:
    pixpts_rot -- namedtuple of numpy arrays of pixel x,y points rotated

    """
    deg2rad = np.pi/180.
    angle_rad = angle*deg2rad
    xpix_pts, ypix_pts = pixpts

    xpix_pts_rotated = xpix_pts*np.cos(angle_rad) - ypix_pts*np.sin(angle_rad)
    ypix_pts_rotated = xpix_pts*np.sin(angle_rad) + ypix_pts*np.cos(angle_rad)

    PixPointsRot = namedtuple('PixPointsRot', 'x y')
    pixpts_rot = PixPointsRot(xpix_pts_rotated, ypix_pts_rotated)

    return pixpts_rot

def translate_pixpts(pixpts_rot, translation, scale_factor=10):
    """
    Geometrically translate rotated pixel points by rover position.

    Keyword arguments:
    pixpts_rot -- namedtuple of numpy arrays of pixel x,y points rotated
    translation -- tuple of displacements along x,y in world frame
    scale_factor -- between world and rover frame pixels

    Return value:
    pixpts_tran -- namedtuple of numpy arrays of pixel x,y points translated

    """
    translation_x, translation_y = translation

    xpix_pts_translated = pixpts_rot.x/scale_factor + translation_x
    ypix_pts_translated = pixpts_rot.y/scale_factor + translation_y

    PixPointsTran = namedtuple('PixPointsTran', 'x y')
    pixpts_tran = PixPointsTran(xpix_pts_translated, ypix_pts_translated)

    return pixpts_tran

def rover_to_world(pixpts_rf, rover_pos, rover_yaw, world_size=200):
    """
    Transform pixel points of ROIs from rover frame to world frame.

    Keyword arguments:
    pixpts_rf -- tuple of numpy arrays of x,y pixel points in rover frame
    rover_pos -- tuple of rover x,y position in world frame
    rover_yaw -- rover yaw angle in world frame
    world_size -- integer length of square world map of 200 x 200 pixels

    Return value:
    pixpts_wf -- namedtuple of numpy arrays of pixel x,y points in world frame

    """
    # Apply rotation and translation
    pixpts_rot = rotate_pixpts(pixpts_rf, rover_yaw)
    pixpts_tran = translate_pixpts(pixpts_rot, rover_pos)

    # Clip pixels to be within world size
    xpix_pts_wf = np.clip(np.int_(pixpts_tran.x), 0, world_size-1)
    ypix_pts_wf = np.clip(np.int_(pixpts_tran.y), 0, world_size-1)

    # Define a named tuple for the points of the three ROIs
    PixPointsWf = namedtuple('PixPointsWf', 'x y')
    pixpts_wf = PixPointsWf(xpix_pts_wf, ypix_pts_wf)

    return pixpts_wf

def inv_translate_pixpts(pixpts_wf, translation, scale_factor=10):
    """
    Inverse translate pixel points from world frame.

    Keyword arguments:
    pixpts_wf -- tuple of numpy arrays of x,y pixel points in world frame
    translation -- tuple of displacements along x,y in world frame
    scale_factor -- between world and rover frame pixels

    Return value:
    pixpts_rot -- namedtuple of numpy arrays of pixel x,y points in prior
                  rotated positions
    """
    translation_x, translation_y = translation
    xpix_pts_wf, ypix_pts_wf = pixpts_wf

    xpix_pts_rotated = (xpix_pts_wf - translation_x)*scale_factor
    ypix_pts_rotated = (ypix_pts_wf - translation_y)*scale_factor

    PixPointsRot = namedtuple('PixPointsRot', 'x y')
    pixpts_rot = PixPointsRot(xpix_pts_rotated, ypix_pts_rotated)

    return pixpts_rot

def inv_rotate_pixpts(pixpts_rot, angle):
    """
    Inverse rotate rotated pixel points to their original positions.

    Keyword arguments:
    pixpts_rot -- namedtuple of numpy arrays of x,y pixel points rotated
    angle -- rotation angle in degrees

    Return value:
    pixpts -- namedtuple of numpy arrays of pixel x,y points in
              original positions
    """
    deg2rad = np.pi/180.
    angle_rad = angle*deg2rad

    xpix_pts = pixpts_rot.x*np.cos(angle_rad) + pixpts_rot.y*np.sin(angle_rad)
    ypix_pts = -pixpts_rot.x*np.sin(angle_rad) + pixpts_rot.y*np.cos(angle_rad)

    PixPoints = namedtuple('PixPoints', 'x y')
    pixpts = PixPoints(xpix_pts, ypix_pts)

    return pixpts

def world_to_rover(pixpts_wf, rover_pos, rover_yaw):
    """
    Transform pixel points of ROIs from world frame to rover frame.

    Keyword arguments:
    pixpts_wf -- tuple of numpy arrays of x,y pixel points in world frame
    rover_pos -- tuple of rover x,y position in world frame
    rover_yaw -- rover yaw angle in world frame

    Return value:
    pixpts_rf -- namedtuple of numpy arrays of pixel x,y points in rover frame

    """
    # Apply inverse translation and rotation
    pixpts_rot = inv_translate_pixpts(pixpts_wf, rover_pos)
    pixpts_rf = inv_rotate_pixpts(pixpts_rot, rover_yaw)

    return pixpts_rf

def perception_step(Rover, R=0, G=1, B=2):
    """
    Sense environment with rover camera and update rover state accordingly.

    Keyword arguments:
    Rover -- instance of RoverTelemetry class
    R,G,B -- indexes representing the RGB color channels in a numpy image

    """
    # Apply perspective transform to get 2D overhead view of rover cam
    warped_img = perspect_transform(Rover.img)

    # Apply color thresholds to extract pixels of navigable/obstacles/rocks
    thresh_pixpts_pf = color_thresh(warped_img)

    # Update rover vision image with each ROI assigned to one of
    # the RGB color channels (to be displayed on left side of sim screen)
    VISION_R_VAL, VISION_G_VAL, VISION_B_VAL = 135, 1, 175
    Rover.vision_image[:, :, R] = thresh_pixpts_pf.obs * VISION_R_VAL
    Rover.vision_image[:, :, G] = thresh_pixpts_pf.rock * VISION_G_VAL
    Rover.vision_image[:, :, B] = thresh_pixpts_pf.nav * VISION_B_VAL

    # Transform pixel coordinates from perspective frame to rover frame
    nav_pixpts_rf = perspect_to_rover(thresh_pixpts_pf.nav)
    obs_pixpts_rf = perspect_to_rover(thresh_pixpts_pf.obs)
    rock_pixpts_rf = perspect_to_rover(thresh_pixpts_pf.rock)

    # Convert above cartesian coordinates to polar coordinates
    Rover.nav_dists, Rover.nav_angles = to_polar_coords(nav_pixpts_rf)
    Rover.obs_dists, Rover.obs_angles = to_polar_coords(obs_pixpts_rf)
    Rover.rock_dists = to_polar_coords(rock_pixpts_rf)[0]

    # Extract subset of nav_angles that are left of rover heading
    Rover.nav_angles_left = Rover.nav_angles[Rover.nav_angles > 0]

    # Only include pixels within certain distances from rover (for fidelity)
    nav_pixpts_rf = [pts[Rover.nav_dists < 60] for pts in nav_pixpts_rf]
    obs_pixpts_rf = [pts[Rover.obs_dists < 80] for pts in obs_pixpts_rf]
    rock_pixpts_rf = [pts[Rover.rock_dists < 70] for pts in rock_pixpts_rf]

    # Convert rock cartesian coords to polar coords
    Rover.rock_angles = to_polar_coords(rock_pixpts_rf)[1]

    # Transform pixel points of ROIs from rover frame to world frame
    nav_pixpts_wf = rover_to_world(nav_pixpts_rf, Rover.pos, Rover.yaw)
    obs_pixpts_wf = rover_to_world(obs_pixpts_rf, Rover.pos, Rover.yaw)
    rock_pixpts_wf = rover_to_world(rock_pixpts_rf, Rover.pos, Rover.yaw)

    # Only update worldmap (displayed on right) if rover has a stable drive
    # High pitch/rolls cause inaccurate 3D to 2D mapping and low fidelity
    is_stable = ((Rover.pitch > 359 or Rover.pitch < 0.25)
                 and (Rover.roll > 359 or Rover.roll < 0.37))

    if is_stable:  # Update map with each ROI assigned to an RGB color channel
        MAP_R_VAL, MAP_G_VAL, MAP_B_VAL = 255, 255, 255
        Rover.worldmap[obs_pixpts_wf.y, obs_pixpts_wf.x, R] += MAP_R_VAL
        Rover.worldmap[rock_pixpts_wf.y, rock_pixpts_wf.x, G] += MAP_G_VAL
        Rover.worldmap[nav_pixpts_wf.y, nav_pixpts_wf.x, B] += MAP_B_VAL

    return Rover


# -*- coding: UTF-8 -*-
"""
A compilation of molecular data and labels as well as a functions to calculate
ionization diagrams and occupation rates.
The variables `H2` and `CO` are dictionaries containing the lines associated
to each molecular band, e.g., AX(1-0) for CO or BX(0-0) for H2 Lyman band.
"""
__author__ = 'Jens-Kristian Krogager'

import numpy as np

CO_full_labels = {
    'AX(0-0)': "${\\rm CO\ A}^1\\Pi(0) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(1-0)': "${\\rm CO\ A}^1\\Pi(1) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(2-0)': "${\\rm CO\ A}^1\\Pi(2) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(3-0)': "${\\rm CO\ A}^1\\Pi(3) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(4-0)': "${\\rm CO\ A}^1\\Pi(4) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(5-0)': "${\\rm CO\ A}^1\\Pi(5) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(6-0)': "${\\rm CO\ A}^1\\Pi(6) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(7-0)': "${\\rm CO\ A}^1\\Pi(7) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(8-0)': "${\\rm CO\ A}^1\\Pi(8) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(9-0)': "${\\rm CO\ A}^1\\Pi(9) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(10-0)': "${\\rm CO\ A}^1\\Pi(10) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'AX(11-0)': "${\\rm CO\ A}^1\\Pi(11) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'CX(0-0)': "${\\rm CO\ C}^1\\Sigma(0) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'dX(5-0)': "${\\rm CO\ d}^3\\Delta(5) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$",
    'eX(1-0)': "${\\rm CO\ e}^3\\Sigma^-(1) \\leftarrow {\\rm X}^1\\Sigma^+(\\nu=0)$"
}

CO_labels = {'COJ0_1544.44': 'AX(0-0)',
             'COJ0_1509.74': 'AX(1-0)',
             'COJ0_1477.56': 'AX(2-0)',
             'COJ0_1447.35': 'AX(3-0)',
             'COJ0_1419.04': 'AX(4-0)',
             'COJ0_1392.52': 'AX(5-0)',
             'COJ0_1367.62': 'AX(6-0)',
             'COJ0_1344.18': 'AX(7-0)',
             'COJ0_1322.15': 'AX(8-0)',
             'COJ0_1301.40': 'AX(9-0)',
             'COJ0_1281.86': 'AX(10-0)',
             'COJ0_1263.43': 'AX(11-0)',
             'COJ0_1087.86': 'CX(0-0)',
             'COJ0_1510.34': 'dX(5-0)',
             'COJ0_1543.17': 'eX(1-0)',
             'COJ0_1543.00': 'eX(1-0)'
             }

H2_full_labels = {
    'BX(0-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(0) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(1-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(1) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(2-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(2) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(3-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(3) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(4-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(4) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(5-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(5) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(6-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(6) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(7-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(7) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(8-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(8) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(9-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(9) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(10-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(10) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(11-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(11) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(12-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(12) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(13-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(13) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(14-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(14) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(15-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(15) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(16-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(16) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(17-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(17) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(18-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(18) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'BX(19-0)': "${\\rm H_2\ B}^1\\Sigma_u^+(19) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'CX(0-0)': "${\\rm H_2\ C}^1\\Pi_u(0) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'CX(1-0)': "${\\rm H_2\ C}^1\\Pi_u(1) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'CX(2-0)': "${\\rm H_2\ C}^1\\Pi_u(2) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'CX(3-0)': "${\\rm H_2\ C}^1\\Pi_u(3) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'CX(4-0)': "${\\rm H_2\ C}^1\\Pi_u(4) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
    'CX(5-0)': "${\\rm H_2\ C}^1\\Pi_u(5) \\leftarrow {\\rm X}^1\\Sigma_g^+(\\nu=0)$",
}

H2_labels = {'H2J0_917.25': 'BX(18-0)',
             'H2J0_1092.20': 'BX(1-0)',
             'H2J0_929.53': 'CX(4-0)',
             'H2J0_946.17': 'BX(14-0)',
             'H2J0_971.99': 'BX(11-0)',
             'H2J0_964.98': 'CX(2-0)',
             'H2J0_1049.37': 'BX(4-0)',
             'H2J0_1012.81': 'BX(7-0)',
             'H2J0_1001.82': 'BX(8-0)',
             'H2J0_981.44': 'BX(10-0)',
             'H2J0_985.63': 'CX(1-0)',
             'H2J0_931.06': 'BX(16-0)',
             'H2J0_1077.14': 'BX(2-0)',
             'H2J0_910.82': 'BX(19-0)',
             'H2J0_938.47': 'BX(15-0)',
             'H2J0_954.41': 'BX(13-0)',
             'H2J0_914.40': 'CX(5-0)',
             'H2J0_1062.88': 'BX(3-0)',
             'H2J0_1036.55': 'BX(5-0)',
             'H2J0_946.43': 'CX(3-0)',
             'H2J0_991.38': 'BX(9-0)',
             'H2J0_1008.55': 'CX(0-0)',
             'H2J0_1108.13': 'BX(0-0)',
             'H2J0_1024.37': 'BX(6-0)',
             'H2J0_923.98': 'BX(17-0)',
             'H2J0_962.98': 'BX(12-0)'
             }

H2 = {'BX(0-0)': [['H2J0_1108.13'],
                  ['H2J1_1110.06', 'H2J1_1108.63'],
                  ['H2J2_1112.50', 'H2J2_1110.12'],
                  ['H2J3_1115.90', 'H2J3_1112.58'],
                  ['H2J4_1120.25', 'H2J4_1116.01'],
                  ['H2J5_1125.54', 'H2J5_1120.40'],
                  ['H2J6_1131.75', 'H2J6_1125.73'],
                  ['H2J7_1151.64']],

      'BX(1-0)': [['H2J0_1092.20'],
                  ['H2J1_1094.05', 'H2J1_1092.73'],
                  ['H2J2_1096.44', 'H2J2_1094.24'],
                  ['H2J3_1099.79', 'H2J3_1096.73'],
                  ['H2J4_1104.08', 'H2J4_1100.16'],
                  ['H2J5_1109.31', 'H2J5_1104.55'],
                  ['H2J6_1115.46', 'H2J6_1109.86'],
                  ['H2J7_1134.90']],

      'BX(2-0)': [['H2J0_1077.14'],
                  ['H2J1_1078.93', 'H2J1_1077.70'],
                  ['H2J2_1081.27', 'H2J2_1079.23'],
                  ['H2J3_1084.56', 'H2J3_1081.71'],
                  ['H2J4_1088.80', 'H2J4_1085.15'],
                  ['H2J5_1093.95', 'H2J5_1089.51'],
                  ['H2J6_1100.02', 'H2J6_1094.80'],
                  ['H2J7_1119.03']],

      'BX(3-0)': [['H2J0_1062.88'],
                  ['H2J1_1064.61', 'H2J1_1063.46'],
                  ['H2J2_1066.90', 'H2J2_1064.99'],
                  ['H2J3_1070.14', 'H2J3_1067.48'],
                  ['H2J4_1074.31', 'H2J4_1070.90'],
                  ['H2J5_1079.40', 'H2J5_1075.24'],
                  ['H2J6_1085.38', 'H2J6_1080.49'],
                  ['H2J7_1103.98']],

      'BX(4-0)': [['H2J0_1049.37'],
                  ['H2J1_1051.03', 'H2J1_1049.96'],
                  ['H2J2_1053.28', 'H2J2_1051.50'],
                  ['H2J3_1056.47', 'H2J3_1053.98'],
                  ['H2J4_1060.58', 'H2J4_1057.38'],
                  ['H2J5_1065.60', 'H2J5_1061.70'],
                  ['H2J6_1071.50', 'H2J6_1066.91'],
                  ['H2J7_1089.71']],

      'BX(5-0)': [['H2J0_1036.55'],
                  ['H2J1_1038.16', 'H2J1_1037.15'],
                  ['H2J2_1040.37', 'H2J2_1038.69'],
                  ['H2J3_1043.50', 'H2J3_1041.16'],
                  ['H2J4_1047.55', 'H2J4_1044.54'],
                  ['H2J5_1052.50', 'H2J5_1048.83'],
                  ['H2J6_1058.32', 'H2J6_1054.00'],
                  ['H2J7_1076.16']],

      'BX(6-0)': [['H2J0_1024.37'],
                  ['H2J1_1025.94', 'H2J1_1024.99'],
                  ['H2J2_1028.11', 'H2J2_1026.53'],
                  ['H2J3_1031.19', 'H2J3_1028.99'],
                  ['H2J4_1035.18', 'H2J4_1032.35'],
                  ['H2J5_1040.06', 'H2J5_1036.60'],
                  ['H2J6_1045.80', 'H2J6_1041.73'],
                  ['H2J7_1063.29']],

      'BX(7-0)': [['H2J0_1012.81'],
                  ['H2J1_1014.33', 'H2J1_1013.44'],
                  ['H2J2_1016.46', 'H2J2_1014.98'],
                  ['H2J3_1019.50', 'H2J3_1017.42'],
                  ['H2J4_1023.44', 'H2J4_1020.77'],
                  ['H2J5_1028.25', 'H2J5_1024.99'],
                  ['H2J6_1033.92', 'H2J6_1030.07'],
                  ['H2J7_1051.07']],

      'BX(8-0)': [['H2J0_1001.82'],
                  ['H2J1_1003.30', 'H2J1_1002.45'],
                  ['H2J2_1005.39', 'H2J2_1003.99'],
                  ['H2J3_1008.39', 'H2J3_1006.41'],
                  ['H2J4_1012.26', 'H2J4_1009.72'],
                  ['H2J5_1017.00', 'H2J5_1013.71'],
                  ['H2J6_1022.59', 'H2J6_1019.02'],
                  ['H2J7_1039.21']],

      'BX(9-0)': [['H2J0_991.38'],
                  ['H2J1_992.81', 'H2J1_992.02'],
                  ['H2J2_994.87', 'H2J2_993.55'],
                  ['H2J3_997.83', 'H2J3_995.97'],
                  ['H2J4_1001.66', 'H2J4_999.27'],
                  ['H2J5_1006.34', 'H2J5_1003.43'],
                  ['H2J6_1011.87', 'H2J6_1008.43'],
                  ['H2J7_1028.41']],

      'BX(10-0)': [['H2J0_981.44'],
                   ['H2J1_982.84', 'H2J1_982.07'],
                   ['H2J2_984.86', 'H2J2_983.59'],
                   ['H2J3_987.77', 'H2J3_985.96'],
                   ['H2J4_991.53', 'H2J4_989.56'],
                   ['H2J5_996.12', 'H2J5_993.49'],
                   ['H2J6_1001.91', 'H2J6_998.43'],
                   ['H2J7_1017.98']],

      'BX(11-0)': [['H2J0_971.99'],
                   ['H2J1_973.34', 'H2J1_972.63'],
                   ['H2J2_975.35', 'H2J2_974.16'],
                   ['H2J3_978.22', 'H2J3_976.55'],
                   ['H2J4_981.95', 'H2J4_979.81'],
                   ['H2J5_986.52', 'H2J5_983.90'],
                   ['H2J6_991.92', 'H2J6_988.81'],
                   ['H2J7_998.11']],

      'BX(12-0)': [['H2J0_962.98'],
                   ['H2J1_964.31', 'H2J1_963.61'],
                   ['H2J2_966.28', 'H2J2_965.05'],
                   ['H2J3_969.09', 'H2J3_967.68'],
                   ['H2J4_972.69', 'H2J4_970.84'],
                   ['H2J5_977.46', 'H2J5_974.89'],
                   ['H2J6_982.73', 'H2J6_979.76'],
                   ['H2J7_988.84']],

      'BX(13-0)': [['H2J0_954.41'],
                   ['H2J1_955.71', 'H2J1_955.07'],
                   ['H2J2_957.65', 'H2J2_956.58'],
                   ['H2J3_960.45', 'H2J3_958.95'],
                   ['H2J4_964.09', 'H2J4_962.15'],
                   ['H2J5_968.56', 'H2J5_966.18'],
                   ['H2J6_973.83', 'H2J6_971.00'],
                   ['H2J7_989.32']],

      'BX(14-0)': [['H2J0_946.17'],
                   ['H2J1_947.51', 'H2J1_946.98'],
                   ['H2J2_949.35', 'H2J2_948.47'],
                   ['H2J3_952.27', 'H2J3_950.82'],
                   ['H2J4_955.85', 'H2J4_954.00'],
                   ['H2J5_960.27', 'H2J5_958.01'],
                   ['H2J6_965.48', 'H2J6_962.82'],
                   ['H2J7_980.76']],

      'BX(15-0)': [['H2J0_938.47'],
                   ['H2J1_939.71', 'H2J1_939.12'],
                   ['H2J2_941.60', 'H2J2_940.63'],
                   ['H2J3_944.33', 'H2J3_942.96'],
                   ['H2J4_947.89', 'H2J4_946.12'],
                   ['H2J5_952.25', 'H2J5_950.07'],
                   ['H2J6_957.41', 'H2J6_954.70'],
                   ['H2J7_972.44']],

      'BX(16-0)': [['H2J0_931.06'],
                   ['H2J1_932.27', 'H2J1_931.73'],
                   ['H2J2_934.14', 'H2J2_933.24'],
                   ['H2J3_936.86', 'H2J3_935.58'],
                   ['H2J4_940.39', 'H2J4_938.73'],
                   ['H2J5_944.72', 'H2J5_942.69'],
                   ['H2J6_949.84', 'H2J6_947.43'],
                   ['H2J7_964.70']],

      'BX(17-0)': [['H2J0_923.98'],
                   ['H2J1_925.17', 'H2J1_924.64'],
                   ['H2J2_927.02', 'H2J2_926.13'],
                   ['H2J3_929.69', 'H2J3_928.44'],
                   ['H2J4_933.17', 'H2J4_931.54'],
                   ['H2J5_937.44', 'H2J5_935.32'],
                   ['H2J6_942.48', 'H2J6_940.49'],
                   ['H2J7_956.99']],

      'BX(18-0)': [['H2J0_917.25'],
                   ['H2J1_918.41', 'H2J1_917.92'],
                   ['H2J2_920.24', 'H2J2_919.42'],
                   ['H2J3_922.89', 'H2J3_921.73'],
                   ['H2J4_926.35', 'H2J4_924.85'],
                   ['H2J5_930.61', 'H2J5_928.76'],
                   ['H2J6_935.63', 'H2J6_933.44'],
                   ['H2J7_950.12']],

      'BX(19-0)': [['H2J0_910.82'],
                   ['H2J1_911.97', 'H2J1_911.48'],
                   ['H2J2_913.77', 'H2J2_912.95'],
                   ['H2J3_916.38', 'H2J3_915.21'],
                   ['H2J4_919.79', 'H2J4_918.15'],
                   ['H2J5_923.96', 'H2J5_922.48'],
                   ['H2J6_928.78', 'H2J6_927.05'],
                   ['H2J7_943.55']],

      'CX(0-0)': [['H2J0_1008.55'],
                  ['H2J1_1009.77', 'H2J1_1008.50'],
                  ['H2J2_1012.17', 'H2J2_1010.94', 'H2J2_1009.02'],
                  ['H2J3_1014.50', 'H2J3_1012.68', 'H2J3_1010.13'],
                  ['H2J4_1017.39', 'H2J4_1014.98', 'H2J4_1011.81'],
                  ['H2J5_1020.80', 'H2J5_1017.83', 'H2J5_1014.24'],
                  ['H2J6_1024.73', 'H2J6_1021.21', 'H2J6_1016.74'],
                  ['H2J7_1039.77', 'H2J7_1035.43']],

      'CX(1-0)': [['H2J0_985.63'],
                  ['H2J1_986.80', 'H2J1_985.64'],
                  ['H2J2_989.09', 'H2J2_987.97', 'H2J2_986.24'],
                  ['H2J3_991.38', 'H2J3_989.73', 'H2J3_987.45'],
                  ['H2J4_994.23', 'H2J4_992.05', 'H2J4_988.87'],
                  ['H2J5_997.64', 'H2J5_994.92', 'H2J5_991.37'],
                  ['H2J6_1001.21', 'H2J6_998.33', 'H2J6_994.26'],
                  ['H2J7_1015.75', 'H2J7_1012.13']],

      'CX(2-0)': [['H2J0_964.98'],
                  ['H2J1_966.10', 'H2J1_965.06'],
                  ['H2J2_968.30', 'H2J2_967.28', 'H2J2_965.80'],
                  ['H2J3_970.56', 'H2J3_969.05', 'H2J3_966.78'],
                  ['H2J4_973.45', 'H2J4_971.39', 'H2J4_968.67'],
                  ['H2J5_976.55', 'H2J5_974.29', 'H2J5_971.07'],
                  ['H2J6_980.50', 'H2J6_977.73', 'H2J6_974.05'],
                  ['H2J7_994.45', 'H2J7_991.16']],

      'CX(3-0)': [['H2J0_946.43'],
                  ['H2J1_947.42', 'H2J1_946.38'],
                  ['H2J2_949.61', 'H2J2_948.62', 'H2J2_947.11'],
                  ['H2J3_951.67', 'H2J3_950.40', 'H2J3_948.42'],
                  ['H2J4_954.47', 'H2J4_952.76', 'H2J4_950.32'],
                  ['H2J5_957.82', 'H2J5_955.68', 'H2J5_952.80'],
                  ['H2J6_961.70', 'H2J6_959.15', 'H2J6_955.98'],
                  ['H2J7_975.30', 'H2J7_972.27']],

      'CX(4-0)': [['H2J0_929.53'],
                  ['H2J1_930.58', 'H2J1_929.69'],
                  ['H2J2_932.60', 'H2J2_931.78', 'H2J2_930.45'],
                  ['H2J3_934.79', 'H2J3_933.58', 'H2J3_931.81'],
                  ['H2J4_937.55', 'H2J4_935.96', 'H2J4_933.79'],
                  ['H2J5_940.88', 'H2J5_938.91', 'H2J5_936.47'],
                  ['H2J6_944.78', 'H2J6_942.42', 'H2J6_939.11'],
                  ['H2J7_958.19', 'H2J7_955.26']],

      'CX(5-0)': [['H2J0_914.40'],
                  ['H2J1_915.40', 'H2J1_914.61'],
                  ['H2J2_917.37', 'H2J2_916.62', 'H2J2_915.43'],
                  ['H2J3_919.54', 'H2J3_918.43', 'H2J3_916.88'],
                  ['H2J4_922.31', 'H2J4_920.83', 'H2J4_919.05'],
                  ['H2J5_925.66', 'H2J5_923.82', 'H2J5_921.22'],
                  ['H2J6_929.70', 'H2J6_927.36', 'H2J6_924.50'],
                  ['H2J7_942.23', 'H2J7_939.98']]
      }

CO = {
    # Nu=0:
    'AX(0-0)': [['COJ0_1544.44'],  # J=0
                ['COJ1_1544.54', 'COJ1_1544.38'],  # J=1
                ['COJ2_1544.72', 'COJ2_1544.57', 'COJ2_1544.34'],  # J=2
                ['COJ3_1544.84', 'COJ3_1544.61', 'COJ3_1544.31'],  # J=3
                ['COJ4_1544.98', 'COJ4_1544.68', 'COJ4_1544.30'],  # J=4
                ['COJ5_1545.14', 'COJ5_1544.76', 'COJ5_1544.31']],   # J=5
    # Nu=1:
    'AX(1-0)': [['COJ0_1509.74'],
                ['COJ1_1509.83', 'COJ1_1509.69'],
                ['COJ2_1510.01', 'COJ2_1509.87', 'COJ2_1509.66'],
                ['COJ3_1510.13', 'COJ3_1509.92', 'COJ3_1509.64'],
                ['COJ4_1510.27', 'COJ4_1509.99', 'COJ4_1509.64']],
    # Nu=2:
    'AX(2-0)': [['COJ0_1477.56'],
                ['COJ1_1477.64', 'COJ1_1477.51'],
                ['COJ2_1477.81', 'COJ2_1477.68', 'COJ2_1477.47'],
                ['COJ3_1477.93', 'COJ3_1477.72', 'COJ3_1477.45'],
                ['COJ4_1478.06', 'COJ4_1477.79', 'COJ4_1477.45']],
    # Nu=3:
    'AX(3-0)': [['COJ0_1447.35'],
                ['COJ1_1447.43', 'COJ1_1447.30'],
                ['COJ2_1447.59', 'COJ2_1447.46', 'COJ2_1447.27'],
                ['COJ3_1447.70', 'COJ3_1447.51', 'COJ3_1447.25'],
                ['COJ4_1447.83', 'COJ4_1447.58', 'COJ4_1447.25']],
    # Nu=4:
    'AX(4-0)': [['COJ0_1419.04'],
                ['COJ1_1419.12', 'COJ1_1419.00'],
                ['COJ2_1419.27', 'COJ2_1419.15', 'COJ2_1418.97'],
                ['COJ3_1419.38', 'COJ3_1419.20', 'COJ3_1418.96'],
                ['COJ4_1419.51', 'COJ4_1419.27', 'COJ4_1418.97']],
    # Nu=5:
    'AX(5-0)': [['COJ0_1392.52'],
                ['COJ1_1392.60', 'COJ1_1392.48'],
                ['COJ2_1392.74', 'COJ2_1392.63', 'COJ2_1392.46'],
                ['COJ3_1392.85', 'COJ3_1392.68', 'COJ3_1392.45'],
                ['COJ4_1392.98', 'COJ4_1392.75', 'COJ4_1392.46']],
    # Nu=6:
    'AX(6-0)': [['COJ0_1367.62'],
                ['COJ1_1367.69', 'COJ1_1367.58'],
                ['COJ2_1367.83', 'COJ2_1367.73', 'COJ2_1367.56'],
                ['COJ3_1367.94', 'COJ3_1367.78', 'COJ3_1367.56'],
                ['COJ4_1368.07', 'COJ4_1367.85', 'COJ4_1367.58'],
                ['COJ5_1368.21', 'COJ5_1367.94', 'COJ5_1367.61']],
    # Nu=7:
    'AX(7-0)': [['COJ0_1344.18'],
                ['COJ1_1344.25', 'COJ1_1344.15'],
                ['COJ2_1344.39', 'COJ2_1344.29', 'COJ2_1344.13'],
                ['COJ3_1344.49', 'COJ3_1344.34', 'COJ3_1344.13'],
                ['COJ4_1344.62', 'COJ4_1344.41', 'COJ4_1344.15'],
                ['COJ5_1344.76', 'COJ5_1344.49', 'COJ5_1344.18']],
    # Nu=8:
    'AX(8-0)': [['COJ0_1322.15'],
                ['COJ1_1322.21', 'COJ1_1322.11'],
                ['COJ2_1322.35', 'COJ2_1322.25', 'COJ2_1322.10'],
                ['COJ3_1322.45', 'COJ3_1322.30', 'COJ3_1322.10'],
                ['COJ4_1322.57', 'COJ4_1322.37', 'COJ4_1322.13'],
                ['COJ5_1322.71', 'COJ5_1322.46', 'COJ5_1322.17']],
    # Nu=9:
    'AX(9-0)': [['COJ0_1301.40'],
                ['COJ1_1301.46', 'COJ1_1301.37'],
                ['COJ2_1301.59', 'COJ2_1301.50', 'COJ2_1301.36'],
                ['COJ3_1301.70', 'COJ3_1301.55', 'COJ3_1301.37'],
                ['COJ4_1301.82', 'COJ4_1301.63', 'COJ4_1301.39'],
                ['COJ5_1301.95', 'COJ5_1301.72', 'COJ5_1301.43']],
    # Nu=10:
    'AX(10-0)': [['COJ0_1281.86'],
                 ['COJ1_1281.92', 'COJ1_1281.83'],
                 ['COJ2_1282.05', 'COJ2_1281.96', 'COJ2_1281.83'],
                 ['COJ3_1282.15', 'COJ3_1282.02', 'COJ3_1281.84'],
                 ['COJ4_1282.27', 'COJ4_1282.09', 'COJ4_1281.84'],
                 ['COJ5_1282.40', 'COJ5_1282.18', 'COJ5_1281.91']],
    # Nu=11:
    'AX(11-0)': [['COJ0_1263.43'],
                 ['COJ1_1263.49', 'COJ1_1263.40'],
                 ['COJ2_1263.61', 'COJ2_1263.53', 'COJ2_1263.40'],
                 ['COJ3_1263.71', 'COJ3_1263.58', 'COJ3_1263.41'],
                 ['COJ4_1263.83', 'COJ4_1263.66', 'COJ4_1263.44'],
                 ['COJ5_1263.96', 'COJ5_1263.75', 'COJ5_1263.49']],

    'eX(1-0)': [['COJ0_1543.17', 'COJ0_1543.00'],
                ['COJ1_1543.20', 'COJ1_1542.91', 'COJ1_1543.17'],
                ['COJ2_1543.26', 'COJ2_1543.44', 'COJ2_1542.85', 'COJ2_1543.27', 'COJ2_1543.23'],
                ['COJ3_1543.35', 'COJ3_1543.66', 'COJ3_1542.83', 'COJ3_1543.37', 'COJ3_1543.33'],
                ['COJ4_1543.48', 'COJ4_1543.90', 'COJ4_1542.83', 'COJ4_1543.49', 'COJ4_1543.45']],

    'CX(0-0)': [['COJ0_1087.86'],
                ['COJ1_1087.95', 'COJ1_1087.82'],
                ['COJ2_1088.00', 'COJ2_1087.77'],
                ['COJ3_1088.04', 'COJ3_1087.72'],
                ['COJ4_1088.09', 'COJ4_1087.67']],

    'dX(5-0)': [['COJ0_1510.34'],
                ['COJ1_1510.42', 'COJ1_1510.30'],
                ['COJ2_1510.60', 'COJ2_1510.48', 'COJ2_1510.29'],
                ['COJ3_1510.74', 'COJ3_1510.56', 'COJ3_1510.31'],
                ['COJ4_1510.91', 'COJ4_1510.66', 'COJ4_1510.36']]
}

# --- Rotatinal Constants in units of cm^-1
#     E = hc * B * J(J + 1)
rotational_constant = {'H2': 60.853,
                       'CO': 1.9313,
                       'HD': 45.655
                       }

# Centrifugal constants in units of cm^-1:
centrifugal_constant = {'H2': 4.71e-2,
                        'CO': 6.12e-6
                        }

hc = 1.2398e-4         # eV.cm
k_B = 8.6173e-5        # eV/K

def energy_of_level(element, J):
    """
    Calculate the energy of a given rotational level, `J`
    for the given molecule with correction for centrigual
    expansion.
    E(J) = B_e * J * (J+1)

    Returns
    =======
    E : float
        The energy of the given level in units of K.
    """
    B_e = rotational_constant[element]
    D_e = centrifugal_constant[element]
    E = B_e * hc/k_B * J * (J+1) - hc*D_e * J**2 * (J+1)**2
    return E

def population_of_level(element, T, J):
    """
    Calculate the population of the Jth level relative to the J=0 level.
    The distribution is assumed to be an isothermal Boltzmann distribution:

    n(J) \\propto g(J) e^(-E(J) / kT)
    """
    if element not in rotational_constant.keys():
        print(" Element is not in database! ")
        print(" All elements in database are: " + ", ".join(rotational_constant.keys()))
        return None

    if element == 'H2':
        def g(J):
            Ij = J % 2
            return (2*J + 1)*(2*Ij + 1)
    elif element == 'CO':
        def g(J):
            return 2*J + 1

    E = energy_of_level(element, J)
    n_J = g(J) * np.exp(-E/T)
    return n_J

def calculate_T(element, logN1, logN2, J1, J2):
    """
    Calculate the isothermal temperature for two given column densities.
    """
    if element == 'H2':
        def g(J):
            Ij = J % 2
            return (2*J + 1)*(2*Ij + 1)
    elif element == 'CO':
        def g(J):
            return 2*J + 1

    E1 = energy_of_level(element, J1)
    E2 = energy_of_level(element, J2)
    E12 = E2-E1
    T = -E12/(np.log(10**(logN2 - logN1) * g(J1)/g(J2)))
    return T


# -*- coding: utf-8 -*-
"""
Forked in Hydra IMF from Hydra/MUSE on Feb 19, 2018

@author: Carlos Eduardo Barbosa

Run pPXF in data
"""
import os
import yaml

import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy import constants
from astropy.table import Table, vstack, hstack

from ppxf import ppxf, ppxf_util

import context
import misc
from der_snr import DER_SNR

def run_ppxf(specs, templates_file, outdir, velscale=None, redo=False,
             degree=-1, mdegree=15, clean=False, V0=None):
    """ Running pPXF. """
    velscale = context.velscale if velscale is None else velscale
    ssp_templates = fits.getdata(templates_file, extname="SSPS").T
    nssps = ssp_templates.shape[1]
    logwave_temp = Table.read(templates_file, hdu=2)["loglam"].data
    wave_temp = np.exp(logwave_temp)
    V0 = context.V if V0 is None else V0
    start0 = [V0, 100., 0., 0.]
    bounds = [[[V0 - 2000., V0 + 2000], [3., 800.]],
               [[V0 - 2000., V0 + 2000], [3., 800.]]]
    if not os.path.exists(outdir):
        os.mkdir(outdir)
    for spec in specs:
        # Reading the data in the files
        name = spec.replace(".fits", "")
        outyaml = os.path.join(outdir, "{}.yaml".format(name))
        if os.path.exists(outyaml) and not redo:
            continue
        print("Processing spectrum {}".format(name))
        table = Table.read(spec)
        wave = table["wave"]
        flux = table["flam"]
        fluxerr = table["flamerr"]
        ###################################################################
        # Trim spectra to conform to the wavelenght range of the templates
        idx = np.argwhere(np.logical_and(
                          wave > wave_temp[0],
                          wave < wave_temp[-1]))
        wave = wave[idx].T[0]
        flux = flux[idx].T[0]
        fluxerr = fluxerr[idx].T[0]
        ####################################################################
        # Rebinning the data to a logarithmic scale for ppxf
        wave_range = [wave[0], wave[-1]]
        galaxy, logLam, vtemp = ppxf_util.log_rebin(wave_range,
                                               flux, velscale=velscale)
        noise = ppxf_util.log_rebin(wave_range, fluxerr,
                               velscale=velscale)[0]
        ####################################################################
        # Setting up the gas templates
        gas_templates, line_names, line_wave = \
            ppxf_util.emission_lines(logwave_temp,
                                     [wave[0], wave[-1]], 2.95)
        ngas = gas_templates.shape[1]
        # Preparing the fit
        start = [start0[:2], start0[:2]]
        dv = (logwave_temp[0] - logLam[0]) * \
             constants.c.to("km/s").value
        templates = np.column_stack((ssp_templates, gas_templates))
        components = np.hstack((np.zeros(nssps), np.ones(ngas))).astype(np.int)
        gas_component = components > 0
        ########################################################################
        # # Masking bad pixels
        # skylines = np.array([4785, 5577, 5889, 6300, 6863])
        # goodpixels = np.arange(len(wave))
        # for line in skylines:
        #     sky = np.argwhere((wave < line - 15) | (wave > line + 15)).ravel()
        #     goodpixels = np.intersect1d(goodpixels, sky)
        # print(len(wave), len(goodpixels))
        # input()
        ###################################################################
        # Fitting with two components
        pp = ppxf.ppxf(templates, galaxy, noise, velscale=velscale,
                  plot=True, moments=[2,2], start=start, vsyst=dv,
                  lam=np.exp(logLam), component=components, degree=degree,
                  gas_component=gas_component, gas_names=line_names,
                  quiet=False, mdegree=mdegree, bounds=bounds, clean=clean)
        plt.savefig(os.path.join(outdir, "{}.png".format(name)), dpi=250)
        plt.close()
        pp.name = name
        pp.fit_sn = float(np.nanmedian(pp.galaxy) / \
                      np.nanstd(pp.galaxy - pp.bestfit))
        pp.der_sn = float(misc.snr(flux)[2])
        # Saving results and plot
        save(pp, outdir)

def save(pp, outdir):
    """ Save results from pPXF into files excluding fitting arrays. """
    array_keys = ["lam", "galaxy", "noise", "bestfit", "gas_bestfit",
                  "mpoly", "apoly"]
    array_keys = [_ for _ in array_keys if isinstance(getattr(pp, _),
                                                      np.ndarray)]
    table = Table([getattr(pp, key) for key in array_keys], names=array_keys)
    table.write(os.path.join(outdir, "{}_bestfit.fits".format(pp.name)),
                overwrite=True)
    ppdict = {}
    save_keys = ["name", "regul", "degree", "mdegree", "reddening", "clean",
                 "ncomp", "chi2", "der_sn", "fit_sn"]
    # Chi2 is a astropy.unit.quantity object, we have to make it a scalar
    pp.chi2 = float(pp.chi2)
    for key in save_keys:
        ppdict[key] = getattr(pp, key)
    klist = ["V", "sigma"]
    for j, sol in enumerate(pp.sol):
        for i in range(len(sol)):
            ppdict["{}_{}".format(klist[i], j)] = float(sol[i])
            ppdict["{}err_{}".format(klist[i], j)] = float(pp.error[j][i])

    with open(os.path.join(outdir, "{}.yaml".format(pp.name)), "w") as f:
        yaml.dump(ppdict, f, default_flow_style=False)

def make_table(direc, output):
    """ Read all yaml files in a ppf directory to one make table for all
    bins. """
    filenames = sorted([_ for _ in os.listdir(direc) if _.endswith(".yaml")])
    keys = ["name", "V_0", "Verr_0", "sigma_0", "sigmaerr_0", "der_sn"]
    names = {"name": "spec", "V_0": "V", "Verr_0": "Verr",
             "sigma_0": "sigma", "sigmaerr_0": "sigmaerr", "der_sn": "SNR"}
    outtable = []
    for fname in filenames:
        with open(os.path.join(direc, fname)) as f:
            props = yaml.load(f)
        data = Table([[props[k]] for k in keys], names=[names[k] for k in keys])
        outtable.append(data)
    outtable = vstack(outtable)
    outtable.write(output, format="fits", overwrite=True)

def run_ngc3311():
    targetSN = 250
    w1 = 4500
    w2 = 10000
    sample = "kinematics"
    velscale = context.velscale
    dataset = "MUSE"
    tempfile = os.path.join(context.home, "templates",
               "emiles_muse_vel{}_w{}_{}_{}_fwhm2.95.fits".format(int(velscale),
                w1, w2, sample))
    fields = context.fields[:1]
    for field in fields:
        wdir = os.path.join(context.get_data_dir(dataset), field,
                            "sn{}/sci".format(targetSN))
        os.chdir(wdir)
        specs = sorted([_ for _ in os.listdir(".") if _.endswith(".fits")])
        outdir = os.path.join(os.path.split(os.getcwd())[0],
                              "ppxf_vel{}_w{}_{}_{}".format(int(velscale),
                                                            w1, w2, sample))
        run_ppxf(specs, tempfile, outdir, redo=False)
        outtable = os.path.join(os.path.split(wdir)[0], \
            "ppxf_vel{}_w{}_{}_{}.fits".format(int(velscale), w1, w2, sample))
        make_table(outdir, outtable)

def run_m87():
    targetSN = 500
    w1 = 4500
    w2 = 10000
    sample = "kinematics"
    velscale = context.velscale
    tempfile = os.path.join(context.home, "templates",
               "emiles_muse_vel{}_w{}_{}_{}_fwhm2.95.fits".format(int(velscale),
                w1, w2, sample))
    imgname, cubename = context.get_img_cube_m87()
    wdir = os.path.join(os.path.split(imgname)[0],
                        "sn{}/sci".format(targetSN))
    os.chdir(wdir)
    specs = sorted([_ for _ in os.listdir(".") if _.endswith(".fits")])
    outdir = os.path.join(os.path.split(os.getcwd())[0],
                          "ppxf_vel{}_w{}_{}_{}".format(int(velscale),
                                                        w1, w2, sample))
    run_ppxf(specs, tempfile, outdir, redo=True, degree=10, mdegree=-1, V0=1284)
    outtable = os.path.join(os.path.split(wdir)[0], \
        "ppxf_vel{}_w{}_{}_{}.fits".format(int(velscale), w1, w2, sample))
    make_table(outdir, outtable)

if __name__ == '__main__':
    # run_ngc3311()
    run_m87()


# The MIT License (MIT)
#
# Copyright (c) 2011, 2013 OpenWorm.
# http://openworm.org
#
# All rights reserved. This program and the accompanying materials
# are made available under the terms of the MIT License
# which accompanies this distribution, and is available at
# http://opensource.org/licenses/MIT
#
# Contributors:
#      OpenWorm - http://openworm.org/people.html
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import with_statement
"""
Simulation of C.Elegans muscle cell with three currents
and one calcium pump.

This is still at a very early stage, being optimized on
CamGrid using neuronoptimizer. The model represented
here requires the publically-unavailable nrntools
module from the nrndev library. Please email
author (vellamike@gmail.com) for these libraries.

Author:Mike Vella
email:mv333@cam.ac.uk

WARNING: Not under development as of 29/8/12. Development focus has shifted
to the pyramidal implementation.
"""

from neuron import h
import neuron
import numpy as np
from matplotlib import pyplot as plt
#from nrndev import nrntools

def set_section_mechanism(sec, mech, mech_attribute, mech_value): #Is this passing by reference, do you need to state that explicitly in Python?
    for seg in sec:
        setattr(getattr(seg, mech), mech_attribute, mech_value)

# Create muscle cell body
muscle=h.Section()

# XXX Placeholder values - check dimensions
# Using dimensions from ../NeuroML2/SingleCompMuscle.cell.nml
# Dimensions are in um
muscle.diam = 5
muscle.L = 20

muscle.insert('na')
muscle.insert('kv')
muscle.insert('canrgc')
muscle.insert('cad2')

# set conductance density
set_section_mechanism(muscle,'canrgc','gbar',0.1)
set_section_mechanism(muscle,'kv','gbar',20)
set_section_mechanism(muscle,'na','gbar',100)

# Loop to connect muscle arms
num_arms = 5 # Number of arms per cell
num_compartments = 10 # Number of compartments per arm
arms = []
for i in range(num_arms):
    arms.append([])
    arm_compartment = None
    prev_compartment = muscle

    # Connect 10 compartments per arm
    for j in range(num_compartments):
        arm_compartment = h.Section()

        # Set dimensions
        arm_compartment.diam = 0.75
        arm_compartment.L = 1

        arm_compartment.cm = 1   # uF/cm2
        arm_compartment.Ra = 100 # Ohm-cm

        # Set currents
        arm_compartment.insert('pas')
        set_section_mechanism(arm_compartment, 'pas', 'g', 0.000022222)

        arm_compartment.connect(prev_compartment)
        arms[i].append(arm_compartment)
        prev_compartment = arm_compartment

    # Set up stimulus at end of each muscle arm
    stim=h.IClamp(arms[i][num_compartments-1](0.5))
    stim.delay=100
    stim.amp=0.02
    stim.dur=200

# Set the recording section - middle of muscle cell:
rec_t=h.Vector()
rec_t.record(h._ref_t)
rec_v=h.Vector()
#rec_v.record(arms[0][9](0.5)._ref_v)
rec_v.record(muscle(0.5)._ref_v)

rec_ina = h.Vector()
rec_ina.record(muscle(0.5)._ref_ina)

rec_ik = h.Vector()
rec_ik.record(muscle(0.5)._ref_ik)

rec_ica = h.Vector()
rec_ica.record(muscle(0.5)._ref_ica)

h.dt=0.05
h.finitialize(-70.0)
neuron.init()
sim_time=500

if sim_time:
    neuron.run(sim_time)
else:
    neuron.run(sim_time)

x=np.array(rec_t)
y=np.array(rec_v)
plt.figure(1)
plt.subplot(411)
plt.plot(x,y)

na=np.array(rec_ina)
plt.subplot(412)
plt.plot(x,na)

k=np.array(rec_ik)
plt.subplot(413)
plt.plot(x,k)

ca=np.array(rec_ica)
plt.subplot(414)
plt.plot(x,ca)

plt.show()

##if using nrntools, much of the above can be accomplished as follows:
#sim=nrntools.Simulation(muscle,sim_time=2200,v_init=-70.0)
#sim.set_IClamp(100, 0.2, 2000)
#sim.go()
#sim.show()


"""
Environment for Intergrated Energy System.
You can customize this script in a way you want.

Requirement:
pyglet >= 1.2.4
numpy >= 1.12.1
"""
import numpy as np

class CHPEnv(object):
    GB_bound = [0.2, 1]
    GT_bound = [0, 1]
    TST_bound = [- 0.1, 0.2]
    Grid_bound = [-1, 1]
    action_bound = [-0.2, 0.2]
    action_dim = 4
    state_dim = 14
    dt = .1  # refresh rate
    get_point = False
    grab_counter = 0

    def __init__(self):
        self.device_info = np.zeros((4, 4))
        self.device_info[0, 0] = 5000
        self.device_info[1, 0] = 5000
        self.device_info[2, 0] = 5000
        self.device_info[3, 0] = 2000
        self.he_q = 0
        self.cost = 0
        self.realcost = 0
        self.total_p = 0
        self.wind = 0
        self.rtp = 0
        self.point_info = np.squeeze(np.array([3500, 8500]))
        self.point_l = 0.075
        self.center = np.squeeze(np.array([2500, 8500]))
        self.point_info_init = self.point_info.copy()
        self.tank_info = 2000

    def step(self, action):
        # action = (node1 angular v, node2 angular v)
        gt_action = np.clip(action[0], *self.action_bound)
        gb_action = np.clip(action[1], *self.action_bound)
        tst_action = np.clip(action[2], *self.action_bound)
        grid_action = np.clip(action[3], *self.action_bound)

        self.device_info[0, 1] += gt_action * self.dt
        self.device_info[0, 1] = np.clip(self.device_info[0, 1], 0.2, 1)
        self.device_info[1, 1] += gb_action * self.dt
        self.device_info[1, 1] = np.clip(self.device_info[1, 1], 0, 1)
        self.device_info[2, 1] += tst_action * self.dt
        self.device_info[2, 1] = np.clip(self.device_info[2, 1], (self.tank_info - 500)/5000, (self.tank_info + 1000)/5000)
        self.device_info[2, 1] = np.clip(self.device_info[2, 1], 0, 1)
        self.device_info[3, 1] += grid_action * self.dt
        self.device_info[3, 1] = np.clip(self.device_info[3, 1], -1, 1)

        gt_p = self.device_info[0, 1]*self.device_info[0, 0]
        self.device_info[0, 2] = gt_p
        gt_q = gt_p*2.3*0.75
        self.device_info[0, 3] = gt_q
        gb_q = self.device_info[1, 1]*self.device_info[1, 0]
        self.device_info[1, 3] = gb_q
        tst_q = np.clip((self.device_info[2, 1]*self.device_info[2, 0] - self.tank_info), -500, 1000)
        self.device_info[2, 3] = tst_q

        # gb_q = self.point_info[1] - gt_q + tst_q
        # self.device_info[1, 1] = np.clip(gb_q / 5000, 0, 1)
        # gb_q = self.device_info[1, 1] * self.device_info[3, 0]
        # buy_p = self.point_info[0] - gt_p - self.wind
        # self.device_info[3, 1] = np.clip(buy_p/2000, -1, 1)
        buy_p = self.device_info[3, 1]*self.device_info[3, 0]
        self.device_info[3, 2] = buy_p

        self.he_q = (gt_q + gb_q - tst_q)  # 换热站总热量
        self.total_p = gt_p + buy_p + self.wind  # 总产生电量

        self.cost = (0.345 * gt_q * 2.3*0.75/(1+2.3*0.75) + 0.345 * gb_q/0.9)/self.point_info[1] +\
                    (buy_p * self.rtp + 0.345 * gt_q * 1/(1+2.3*0.75))/self.point_info[0]

        self.realcost = 0.345 * gt_q * 2.3*0.75 + 0.345 * gb_q/0.9 + buy_p * self.rtp

        s, p_distance, q_distance = self._get_state()
        r = self._r_func(p_distance, q_distance, tst_q)

        return s, r, self.get_point

    def reset(self):
        self.get_point = False
        self.grab_counter = 0
        # price = np.array([0.427, 0.427, 0.427, 0.427, 0.427, 0.427, 0.527, 0.527, 0.627, 0.627, 0.627,
        #                 0.527, 0.527, 0.527, 0.527, 0.527, 0.527, 0.627, 0.627, 0.627, 0.627, 0.627,
        #                 0.427, 0.427])
        # wind_power = np.array([875, 1234, 1390, 1392, 1336, 1223, 1173, 1136, 1158, 1312, 1369,
        #                       1376, 1315, 1301, 1343, 1310, 1208, 1055, 896, 773, 672, 626,
        #                       624, 703])

        # temperature = 2*np.array([4800, 4896, 4953.6, 4992, 4896, 4800, 4560, 4320, 4128,
        #                           3984, 3888, 3801.6, 3758.4, 3744, 3748.8, 3772.8, 3820.8,
        #                           3888, 4032, 4224, 4416, 4608, 4752, 4800])
#
        # e_load = np.array([2178.0, 2009, 1873, 1755, 1704, 1839, 2517, 4211, 5397, 5375,
        #                    5651, 5481, 5227, 5176, 5143, 5227, 5909, 6417, 6545, 6206,
        #                   5698, 4510, 4025, 2093])

        p = np.random.uniform(1500, 7000)
        q = np.random.uniform(5000, 11000)
        t = np.random.uniform(200, 2000)
        rtp = np.random.uniform(0.4, 0.7)
        tank = np.random.uniform(1000, 3000)
        self.point_info[0] = p
        self.point_info[1] = q
        self.wind = t
        self.rtp = rtp
        self.tank_info = tank
        return self._get_state()[0]

    def set(self):
        self.get_point = False
        self.grab_counter = 0
        self.point_info[0] = 2000
        self.point_info[1] = 9000
        self.wind = 700
        self.rtp = 0.627
        self.tank_info = 1000
        return self._get_state()[0]

    def render(self):
        s = self._get_state()
        print(s, self.realcost)

    def sample_action(self):
        return np.random.uniform(*self.action_bound, size=self.action_dim)

    def _get_state(self):
        # return the distance (dx, dy) between arm finger point with blue point
        he_q = self.he_q
        total_p = self.total_p
        t_p = total_p - self.point_info[0]
        t_q = he_q - self.point_info[1]
        # dis = self.point_info[1]/self.point_info[0]
        # ratio = self.device_info[0, 1] - self.device_info[1, 1]
        p_gt = self.device_info[0, 2]/self.point_info[0]
        q_gt = self.device_info[0, 3]/self.point_info[1]
        q_gb = self.device_info[1, 3]/self.point_info[1]
        q_tst = self.device_info[2, 3]/self.point_info[1]
        p_grid = self.device_info[3, 2]/self.point_info[0]
        p_wind = self.wind/self.point_info[0]
        cen_dis_p = (self.center[0] - self.point_info[0])/5000
        cen_dis_q = (self.center[1] - self.point_info[1])/5000
        in_point = 1 if self.grab_counter > 0 else 0
        return np.hstack([in_point, t_p/5000,  t_q/5000, cen_dis_p, cen_dis_q, p_gt, q_gt,
                          q_gb, q_tst, p_grid, p_wind, self.device_info[2, 1], self.tank_info/5000, self.rtp
                          # arm1_distance_p, arm1_distance_b,
                          ]), t_p/5000, t_q/5000

    def _r_func(self, p_distance, q_distance, tst_q):
        t = 30
        abs_distance = np.sqrt(np.square(p_distance)+np.square(q_distance))
        r = - abs_distance - self.cost - 0.1*np.abs(self.device_info[2, 1] - 0.4)
        # print(0.1*np.abs(self.device_info[2, 1] - 0.4))
        # print(-self.cost)
        # print(tst_q*(self.rtp-self.rtp_m)/1000)
        if abs_distance < self.point_l and (not self.get_point):
            r += 1
            self.grab_counter += 1
            if self.grab_counter > t:
                r += 10.
                self.get_point = True
        elif abs_distance > self.point_l:
            self.grab_counter = 0
            self.get_point = False
        return r

if __name__ == '__main__':
    env = CHPEnv()
    action = env.sample_action()
    print(action)
    env.reset()
    print(env.tank_info)
    s, r, env.get_point = env.step(action)
    print(env.device_info)


#!/usr/bin/env python3
#PYTHON_ARGCOMPLETE_OK

import logging
import itertools
from dataclasses import dataclass
from simple_parsing import Serializable
from typing import Union, Tuple

import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F
from torchvision.utils import save_image

try:
    from pytorch3d.structures import Pointclouds, Meshes
    from pytorch3d.transforms.transform3d import Transform3d
    from pytorch3d.renderer import (
        look_at_view_transform,
        FoVPerspectiveCameras,
        PointsRasterizationSettings,
        PointsRenderer,
        PointsRasterizer,
        RasterizationSettings,
        MeshRasterizer,
        PointLights,
        MeshRenderer,
        SoftPhongShader,
        AlphaCompositor,
        TexturesVertex
    )
except ImportError:
    logging.warn(
        'pytorch3d import has failed, colored cube dataset will be disabled.')
from top.run.torch_util import resolve_device
from top.run.app_util import update_settings
from top.data.schema import Schema

class ColoredCubeDataset(th.utils.data.IterableDataset):
    """Toy generative dataset for 3D object detection: vertices of an oriented
    unit cube rendered as a point cloud. Intended to be an "easy" baseline.

    # FIXME(ycho): I realized that the cube length is actually 2
    """

    @dataclass
    class Settings(Serializable):
        batch_size: int = 1
        aspect: float = 1.0  # pixel aspect ratio, max_x/max_y
        fov: float = 60  # full vertical field of view, in degrees.
        znear: float = 0.1
        zfar: float = 100.0
        min_distance: float = 0.1
        max_distance: float = 10.0
        image_size: Tuple[int, int] = (256, 256)  # Order: H W
        # Unstack output tensors, for compatibility with Objectron.
        unstack: bool = True
        use_mesh: bool = False

    def __init__(self, opts: Settings, device: th.device = '', transform=None):
        super().__init__()
        self.opts = opts
        self.device = resolve_device(device)
        self.xfm = transform
        # TODO(ycho): Consider support for multiple "objects".
        # TODO(ycho): Consider support for *animated* objects through time.

        self.cloud = self._get_cube_cloud()
        self.clouds = self.cloud.extend(self.opts.batch_size)
        self.mesh = self._get_cube_mesh()
        self.meshes = self.mesh.extend(self.opts.batch_size)

        self.renderer = self._setup_render()
        self.tan_half_fov = np.tan(np.deg2rad(0.5 * self.opts.fov))

        self.min_distance = np.maximum(
            opts.min_distance,
            (0.5 * np.sqrt(3)) / self.tan_half_fov)

    def _setup_render(self):
        # Unpack options ...
        opts = self.opts

        # Initialize a camera.
        # TODO(ycho): Alternatively, specify the intrinsic matrix `K` instead.
        cameras = FoVPerspectiveCameras(
            znear=opts.znear,
            zfar=opts.zfar,
            aspect_ratio=opts.aspect,
            fov=opts.fov,
            degrees=True,
            device=self.device
        )

        # Define the settings for rasterization and shading.
        # As we are rendering images for visualization purposes only we will set faces_per_pixel=1
        # and blur_radius=0.0. Refer to raster_points.py for explanations of
        # these parameters.
        # points_per_pixel (Optional): We will keep track of this many points per
        # pixel, returning the nearest points_per_pixel points along the z-axis

        # Create a points renderer by compositing points using an alpha compositor (nearer points
        # are weighted more heavily). See [1] for an explanation.
        if self.opts.use_mesh:
            raster_settings = RasterizationSettings(
                image_size=opts.image_size,
                blur_radius=0.0,  # hmm...
                faces_per_pixel=1
            )
            rasterizer = MeshRasterizer(cameras=cameras,
                                        raster_settings=raster_settings)
            lights = PointLights(device=self.device,
                                 location=[[0.0, 0.0, -3.0]])

            renderer = MeshRenderer(
                rasterizer=rasterizer,
                shader=SoftPhongShader(
                    device=self.device,
                    cameras=cameras,
                    lights=lights
                )
            )
        else:
            raster_settings = PointsRasterizationSettings(
                image_size=opts.image_size,
                radius=0.1,
                points_per_pixel=8
            )
            rasterizer = PointsRasterizer(
                cameras=cameras, raster_settings=raster_settings)
            renderer = PointsRenderer(
                rasterizer=rasterizer,
                compositor=AlphaCompositor()
            )
        return renderer

    def _get_cube_cloud(self):
        """Get vertices of a unit-cube, with colors assigned according to
        vertex coordinates."""
        vertices = list(itertools.product(
            *zip([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5])))
        vertices = np.insert(vertices, 0, [0, 0, 0], axis=0)
        vertices = th.as_tensor(vertices, dtype=th.float32, device=self.device)

        # Map vertices to colors. =RGB(0.25~0.75)
        colors = (0.5 + 0.5 * vertices)
        cloud = Pointclouds(points=vertices[None], features=colors[None])
        return cloud

    def _get_cube_mesh(self):
        # NOTE(ycho): duplicated from _get_cube_cloud()
        vertices = list(itertools.product(
            *zip([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5])))
        vertices = np.insert(vertices, 0, [0, 0, 0], axis=0)
        vertices = th.as_tensor(
            vertices, dtype=th.float32, device=self.device)

        # FIXME(ycho): Hardcoded face indices.
        # (We can't use Box.FACE since we need triangulated faces)
        faces = [[7, 3, 5],
                 [5, 3, 1],
                 [5, 1, 6],
                 [6, 1, 2],
                 [6, 2, 8],
                 [8, 2, 4],
                 [8, 4, 7],
                 [7, 4, 3],
                 [3, 4, 1],
                 [1, 4, 2],
                 [8, 7, 6],
                 [6, 7, 5]]
        face_indices = th.as_tensor(
            faces, dtype=th.int32, device=self.device).reshape(-1, 3, 3)
        textures = TexturesVertex(
            verts_features=(
                0.5 + 0.5 * vertices)[None])
        mesh = Meshes(
            verts=vertices[None],
            faces=face_indices[None],
            textures=textures
        )

        return mesh

    def _render(self):
        """Render the unit cube at various camera poses.

        We deliberately sample the camera poses such that all vertices
        of the cube are always in view.
        """
        opts = self.opts

        # TODO(ycho): is `no_grad()` here necessary?
        with th.no_grad():
            # Sample a Unit-ray viewing direction.
            ray = th.randn(size=(opts.batch_size, 3), device=self.device)
            ray /= th.norm(ray, dim=1, keepdim=True)

            # Compute distance along the ray according to constraints.
            distance = (self.min_distance +
                        (opts.max_distance - self.min_distance) *
                        th.rand(size=(opts.batch_size, 1), device=self.device))

            # NOTE(ycho): `sqrt(3)/2` here comes from max radius of unit cube.
            # The generic thing to do would be to compute the radius of our
            # cloud. We're not explicitly taking the cube into account.
            max_tangential_offset = th.clamp(
                distance *
                self.tan_half_fov -
                (0.5 * np.sqrt(3)),
                min=0.0)

            # NOTE(ycho): We're sampling an offset orthogonal to the view ray.
            # The constraint here is to be visible within the view plane.
            offset = th.randn(size=(opts.batch_size, 3), device=self.device)
            offset = offset - th.einsum('...a,...a->...',
                                        offset, ray)[:, None] * ray
            offset *= (max_tangential_offset /
                       th.norm(offset, dim=1, keepdim=True))

            # Finally, we have the full definition of the ray.
            at = offset
            pos = offset + ray * distance

            # Compose the `lookat` camera transform according to this position
            R, T = look_at_view_transform(
                eye=pos, at=at, device=self.device)
            if self.opts.use_mesh:
                # NOTE(ycho): Ignoring alpha dimension
                img = self.renderer(
                    self.meshes, R=R, T=T
                )[..., :3]
            else:
                img = self.renderer(
                    point_clouds=self.clouds, R=R, T=T
                )

            # pytorch3d uses `NHWC` convension, convert to NCHW.
            img = img.permute(0, 3, 1, 2)

            # pytorch3d uses a `float` tensor for representing images,
            # whereas we'd like for the image to come out as `uint8` ONLY for
            # compatibility with the output from the `Objectron` dataset.
            img = img.mul_(255.0).to(dtype=th.uint8)

        return (img, R, T)

    def __iter__(self):
        while True:
            (img, R, T) = self._render()
            points_2d = self.renderer.rasterizer.cameras.transform_points(
                self.clouds.points_padded(), R=R, T=T)

            # Convert `points` to UV coordinates.
            # This is to be consistent with the objectron dataset convention.
            if not isinstance(points_2d, th.Tensor):
                points_2d = points_2d.points_padded()
            points_2d[..., :2] *= -0.5
            points_2d[..., :2] += 0.5

            # Add an axis indicating num_instance == 0
            points_2d = points_2d[:, None]

            # Projection matrix maps to -1 ~ 1 NDC coordinates.
            P = self.renderer.rasterizer.cameras.get_projection_transform()
            P = P.get_matrix()

            # NOTE(ycho): Permute projection matrix to fit objectron
            # convention.
            # TODO(ycho): Verify if this keeps +z-axis sign convention
            # consistent.
            permutation = th.as_tensor([
                [0, 1, 0, 0],
                [1, 0, 0, 0],
                [0, 0, -1, 0],
                [0, 0, 0, 1]], dtype=th.float32, device=self.device)
            projection = permutation @ P.transpose(2, 1)

            # NOTE(ycho): Swap conventions here...
            # rotation matrix : rmul -> lmul
            # NOTE(ycho): Also, return flattened output as in `Objectron`.
            translation = T.reshape(self.opts.batch_size, 1, 3)
            scale = th.full_like(translation, 1.0)
            orientation = (
                R.transpose(
                    2, 1).reshape(
                    self.opts.batch_size, 1, -1))
            print('scale', scale.shape)

            # TODO(ycho): Figure out a way to unify these formats.
            # see ai604-video-object-pose#10
            out = {
                Schema.IMAGE: img,
                # NOTE(ycho): For now, only have `1` object per image.
                Schema.CLASS: th.zeros((self.opts.batch_size, 1), dtype=th.int32,
                                       device=self.device),
                Schema.ORIENTATION: orientation,
                Schema.TRANSLATION: translation,
                Schema.SCALE: scale,
                Schema.KEYPOINT_2D: points_2d,
                Schema.INSTANCE_NUM: th.ones(
                    self.opts.batch_size,
                    device=self.device),
                Schema.PROJECTION: projection.repeat(self.opts.batch_size, 1, 1),
            }

            # Unstack the batched render into a set of images, for compatibility with Objectron.
            # NOTE(ycho): See if this results in a significant performance hit.
            if self.opts.unstack:
                for i in range(self.opts.batch_size):
                    out_i = {k: v[i] for k, v in out.items()}
                    if self.xfm is not None:
                        out_i = self.xfm(out_i)
                    yield out_i
            else:
                if self.xfm is not None:
                    out = self.xfm(out)
                yield out

def main():
    opts = ColoredCubeDataset.Settings(batch_size=32, unstack=False)
    opts = update_settings(opts)

    device = resolve_device()
    dataset = ColoredCubeDataset(opts, device,)
    for data in dataset:
        print({k: v.shape for k, v in data.items()})
        save_image(data[Schema.IMAGE] / 255.0, F'/tmp/img.png')
        break

if __name__ == '__main__':
    main()