import sys
import csv
import collections
import math
import numpy as np
import pandas as pd

def delete_false_co(codict,mdict):
    '''Remove false positive crossover events
    based on nTypedBetween and return filtered
    dict.
    '''
    print("Crossovers before filtering: " + str(len(codict)))
    # Use set as there can be quadruple crossovers
    # This leads to duplicate bad keys
    bad_keys = set()
    for key, value in codict.items():
        #nTypedBetween
        ntb = value[9]
        ind = value[1]
        genpos = value[2]
        # Skip last CO for individual
        # These can't be used for filtering
        if ntb != "NA":
            chr_pos1 = value[0] + "_" + value[3]
            co_index1 = codict.keys().index(key)
            # Get following crossover entry
            co_index2 = co_index1 + 1
            co_key2 = list(codict)[co_index2]
            chr_pos2 = codict[co_key2][0] + "_" + codict[co_key2][3]

            #Physical marker positions
            phypos1 = mdict[chr_pos1].split("_")[1]
            phypos2 = mdict[chr_pos2].split("_")[1]
            phychr1 = mdict[chr_pos1].split("_")[0]
            phychr2 = mdict[chr_pos2].split("_")[0]
            if phychr1 == phychr2:
                phydist = int(phypos2) - int(phypos1)
                # Use absolute numbers
                phydist = abs(phydist)
            else:
                # If adjacent marker is from different chromosome
                # Set arbitrary high physical distance
                phydist = 1000000
            # Genetic distance between breakpoints
            gendist = float(codict[co_key2][2]) - float(genpos)
           # Recombinations separated by less than 5 markers are false
           # Recombinations <10kb are considered gene conversions
            if int(ntb) < 5 or phydist < 10000 or gendist < 0.02:
                co_index1 = codict.keys().index(key)
                # Get following crossover entry
                co_index2 = co_index1 + 1
                co_key2 = list(codict)[co_index2]
                bad_keys.add(key)
                if ind == codict[co_key2][1]:
                    bad_keys.add(co_key2)
                else:
                    print("Aborted. Expected paired false positives.")
                    break
    # delete all the bad entries
    for bkey in bad_keys:
        codict.pop(bkey, None)
    print("Crossovers after filtering: " + str(len(codict)))
    return codict

# Read marker positions and crossover into dicts
# These tables are generated with r/qtl locateXO and a custom script

crossovers = sys.argv[1]
# Example of head -2 of input file:
#chr,IND,location,left,right,ileft,iright,gleft,gright,nTypedBetween
#A01,1,111.50960844248,109.290254466964,113.728962417996,402,403,3,2,323

markers = sys.argv[2]
# Example of head -3 of input file:
#chr,name,value
#A01,A01_3072290,0
#A01,A01_3071845,1.35449890934482

# Values in the markers file are genetic map positions in cM

with open(crossovers, mode='r') as co:
    reader1 = csv.reader(co)
    next(reader1, None)
    # Use chr_ind_pos as unique key and all row as values
    codict = collections.OrderedDict((row[0] + "_" + row[1] + "_" + row[2], row[0:]) for row in reader1)
with open(markers, mode='r') as mark:
    reader2 = csv.reader(mark)
    next(reader2, None)
    mdict = collections.OrderedDict((row[0] + "_" + row[2], row[1]) for row in reader2)

# Filter false positive CO
filtdict = delete_false_co(codict,mdict)
# Physical distances between CO
ind_list = []
co = {}
for key, value in filtdict.items():
    ind = value[1]
    chrom = value[0]
    lchrpos = value[0] + "_" + value[3]
    lphypos = mdict[lchrpos].split("_")[1]
    lphychr = mdict[lchrpos].split("_")[0]
    rchrpos = value[0] + "_" + value[4]
    rphypos = mdict[rchrpos].split("_")[1]
    rphychr = mdict[rchrpos].split("_")[0]

    if rphychr == chrom and lphychr == chrom:
        breakpoint = int((int(lphypos) + int(rphypos)) / 2)
        if chrom in co:
            # append CO to list
            breaklist = co[chrom]
            breaklist.append(breakpoint)
            co[chrom] = breaklist
        else:
            co[chrom] = [breakpoint]

    ind_list.append(ind)

# CO counts per individual
counter=collections.Counter(ind_list)
ind_list = []
frq = []
for ind in counter.most_common():
    ind_list.append(ind[0])
    frq.append(ind[1])
inddf = pd.DataFrame({'Ind' : ind_list, 'CO' : frq})
inddf.sort_values(['Ind'], ascending=True, axis=0,inplace=True)
inddf.to_csv('Individual_CO_Frequency.csv', index=False, encoding='utf-8', columns=['Ind','CO'])

# Histogram of CO across chromsomes
chromlist = []
histlist = []
binlist = []

for key, value in co.items():
    chrom = key
    breakpoints = np.array(value)
    # Round highest position up to nearest 500kb
    toprange = math.ceil(breakpoints.max() / 500000.0) * 500000.0
    # Number of bins
    binnum = int(toprange / 500000.0)
    # Calculate histogram
    hist,bins = np.histogram(breakpoints,bins=binnum,range=(0, toprange))
    histlist.extend(np.ndarray.tolist(hist))
    bins_tidy = np.ndarray.tolist(bins)[:-1]
    binlist.extend(bins_tidy)
    for bin in bins_tidy:
        chromlist.append(chrom)
codf = pd.DataFrame({'Chr' : chromlist, 'Bin' : binlist, 'Histogram' : histlist})
codf.sort_values(['Chr', 'Bin'], ascending=True, axis=0,inplace=True)
codf.to_csv('CO_hist.csv', index=False, encoding='utf-8', columns=['Chr','Bin','Histogram'])


#!/usr/bin/python3

"""Projection class
Implements a class for 2D phantoms. The Golosio phantom is coded in terms of
geometry and composition data, with class methods supplied to instantiate it.
Alternatively, phantoms can be defined as a pair of files, one containing
geometry data as a greyscale bitmap and the other containing the composition
data structure defined in a yaml file.

"""

from __future__ import absolute_import, division, print_function
from . import config
import logging
import os
import glob
import fnmatch
from progressbar import ProgressBar

import numpy as np
from numpy import exp, pi
import scipy.constants as sc
from scipy.special import expm1
import matplotlib.pyplot as plt
import xraylib as xrl

from . import helpers
from .helpers import write_tiff32, zero_outside_circle, rotate, imshow
from .data_helpers import MatrixProperties

logger = logging.getLogger(__name__)

UM_PER_CM = 1e4
J_PER_KEV = 1e3 * sc.eV
deg_to_rad = lambda x: x / 180. * pi
rad_to_deg = lambda x: x * 180. / pi

# maia_d = Maia()  # Maia detector object (should be a singleton)

def absorption_sinogram(p, anglelist):
    """Generates the absorption sinogram for absorption by the full
    elemental content of the Phantom2d object.

    Parameters
    ----------
    p : Phantom2d object
    anglelist : list of float
        Ordered list of sinogram projection angles in degrees.

    Returns
    -------
    array of float
        Sinogram of requested scattering or fluorescence.
        This is a 2d x-theta map of dimensionless values.

    """
    sinogram = np.empty((p.cols, len(anglelist)))
    if config.show_progress:
        pbar = ProgressBar(maxval=max(1, len(anglelist)-1), term_width=80).start()
    for i, angle in enumerate(anglelist):
        if config.show_progress:
            pbar.update(i)

        increasing_ix = True   # Set True to accumulate cmam along increasing y
        n_map = irradiance_map(p, angle, n0=1.0, increasing_ix=increasing_ix)
        if increasing_ix:
            sinogram[:, i] = np.log(n_map[0] / n_map[-1])
        else:
            sinogram[:, i] = np.log(n_map[-1] / n_map[0])
    return sinogram

def outgoing_cmam(p, q, maia_d, angle, energy, increasing_ix=True):
    """Compute and return the outgoing cumulative multiplicative absorption map
    (cmam), aka xi'.

    Parameters
    ----------
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    q : int
        Maia detector channel id
    maia_d : Maia() instance
    angle : float
        Tomography projection angle (degree)
    energy : float
        outgoing radiation energy (keV)
    increasing_ix : bool, optional
        If False, performs cumulative sum in opposite direction (in direction of
        decreasing y-index). Note, for the standard Radon transform,
        this direction is unimportant, but for the Radon transform with
        attenuation, we should project in the beam propagation direction.
        (default True).

    Returns
    -------
    2d ndarray of float
        cumulative multiplicative absorption map.

    """
    # The linear absorption map mu = ma_M * mu_M + sum_k ( ma_k * mu_k )
    mu = p.matrix.ma(energy) * p.el_maps['matrix']
    for el in p.el_maps:
        if el == 'matrix':
            continue
        Z = xrl.SymbolToAtomicNumber(el)
        mu += xrl.CS_Total(Z, energy) * p.el_maps[el]

    # project at angle theta by rotating the phantom by angle -theta and
    # projecting along the z-axis (along columns)

    # rotate by the sum of the projection angle and local detector angle
    # Get angle to rotate maps so that propagation toward detector plane
    # corresponds with direction to maia detector element
    phi_x = maia_d.pads[q].angle_X_rad
    phix_deg = rad_to_deg(phi_x)

    # Apply local rotation, accumulation and rotation operators

    # Apply R_{-theta} operator to mu
    # Apply R_{-phi} operator
    mu = rotate(mu, angle + 180 - phix_deg)

    saveflag = np.isclose(angle, config.cmam_angle_save) and config.save_projection_images
    if saveflag:
        path = os.path.join(config.mlem_im_path, 'project_cmam1_%03d.tif' % config.i)
        helpers.write_tiff32(path, mu)

    # Apply C_z operator
    if increasing_ix:
        mu = np.cumsum(mu, axis=0)
    else:
        mu = np.cumsum(mu[::-1], axis=0)[::-1]

    if saveflag:
        path = os.path.join(config.mlem_im_path, 'project_cmam2_%03d.tif' % config.i)
        helpers.write_tiff32(path, mu)

    # Apply R_{phi} operator
    mu = rotate(mu, 180 + phix_deg)
    t = p.um_per_px / UM_PER_CM

    if saveflag:
        path = os.path.join(config.mlem_im_path, 'project_cmam3_%03d.tif' % config.i)
        helpers.write_tiff32(path, mu)

    phi = maia_d.pads[q].out_of_xz_plane_angle
    cmam = mu * t / np.cos(phi)
    return cmam

def project_sinogram(event_type, p, q, maia_d, anglelist, el=None):
    """Generates the sinogram of the requested element accounting for
    absorption by the Phantom2d composition and the geometry.

    Parameters
    ----------
    event_type : string
        One of ['rayleigh', 'compton', 'fluoro'].
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    q : int
        Maia detector channel id
    maia_d : Maia() instance
    anglelist : list of float
        Ordered list of sinogram projection angles in degrees.
    el : string, optional
        Name of element (e.g. 'Fe') used if projecting that element's
        fluorescence.

    Returns
    -------
    array of float
        Sinogram of requested scattering or fluorescence.

    """
    assert event_type in ['rayleigh', 'compton', 'fluoro']

    sinogram = np.empty((p.cols, len(anglelist)))
    if config.show_progress:
        if len(anglelist) > 1:
            pbar = ProgressBar(maxval=len(anglelist)-1).start()
        else:
            pbar = ProgressBar(maxval=1).start()
    for i, angle in enumerate(anglelist):
        if config.show_progress:
            pbar.update(i)

        increasing_ix = True   # Set True to accumulate cmam along increasing y
        n_map = irradiance_map(p, angle, n0=1.0, increasing_ix=increasing_ix)
        e_map = channel_fluoro_map(p, q, maia_d, n_map, angle, el)
        energy = outgoing_photon_energy(event_type, p, q, maia_d, el)
        c = outgoing_cmam(p, q, maia_d, angle, energy, increasing_ix=increasing_ix)
        if config.no_out_absorption:
            sinogram[:, i] = e_map.sum(axis=0)
        else:
            sinogram[:, i] = (e_map * np.exp(-c)).sum(axis=0)

        if i == config.angle_save_index and config.save_projection_images:
            path = os.path.join(config.mlem_im_path, 'project_%s_map_%03d.tif' % (el, config.i))
            helpers.write_tiff32(path, p.el_maps[el])
            path = os.path.join(config.mlem_im_path, 'project_n_map_%03d.tif' % config.i)
            helpers.write_tiff32(path, n_map)
            path = os.path.join(config.mlem_im_path, 'project_e_map_%03d.tif' % config.i)
            helpers.write_tiff32(path, e_map)
            path = os.path.join(config.mlem_im_path, 'project_c_map_%03d.tif' % config.i)
            helpers.write_tiff32(path, c)

    return sinogram

def irradiance_map(p, angle, n0=1.0, increasing_ix=True, matrix_only=False):
    """Generates the image-sized map of irradiance [1/(cm2 s)] at each 2d pixel
    for a given angle accounting for absorption at the incident energy by the
    full elemental distribution.
    Note: In the full biological approximation, where we only need to consider absorption by
    the matrix and not by the full elemental composition, we could cache this result,
    possibly by adding the memoization decorator I use in Sakura here:
    https://github.com/AustralianSynchrotron/Sakura/blob/master/utils.py
    https://github.com/AustralianSynchrotron/Sakura/blob/master/memoize_core.py

    Parameters
    ----------
    p : phantom object
        p.energy - incident beam energy (keV)
        p.um_per_px - length of one pixel of the map (um)
    angle : float
        tomography angle in degrees. The rotation axis is the y-axis in a conventional xyz
        right-handed coordinate system, so positive rotation in the xz-plane is ccw. i.e.
        with y(out-of-page) o--> x
                            |
                            v
                            z
        therefore, using our rotate function, which assumes 2d rotation
    n0 : float
        incident irradiance (default 1.0).
    increasing_ix : bool, optional
        If False, performs cumulative sum in opposite direction (in direction of
        decreasing y-index). Note, for the standard Radon transform,
        this direction is unimportant, but for the Radon transform with
        attenuation, we should project in the beam propagation direction.
        (default True).
    matrix_only : bool, optional
        If True, the map only considers the matrix and doesn't consider the
        other elements in the model.
        (default False).

    Returns
    -------
    2d ndarray of float
        The irradiance map.

    """
    # matrix_map = zero_outside_circle(p.el_maps['matrix'])

    # The linear absorption map mu0 = ma_M * mu_M + sum_k ( ma_k * mu_k )
    mu0 = p.matrix.ma(p.energy) * p.el_maps['matrix']
    if not config.absorb_with_matrix_only:
        for el in p.el_maps:
            if el == 'matrix':
                continue
            Z = xrl.SymbolToAtomicNumber(el)
            mu0 += xrl.CS_Total(Z, p.energy) * p.el_maps[el]

    # project at angle theta by rotating the phantom by angle -theta and
    # projecting along the z-axis (along columns)
    im = rotate(mu0, angle)     # rotate by angle degrees ccw
    t = p.um_per_px / UM_PER_CM
    # accumulate along z-axis (image defined in xz-coordinates, so accumulate
    # along image rows), consistent with matlab sinogram convention.
    # See http://www.mathworks.com/help/images/radon-transform.html
    if increasing_ix:
        cmam = t * np.cumsum(im, axis=0)
    else:
        cmam = t * np.cumsum(im[::-1], axis=0)[::-1]
    if config.no_in_absorption:
        n_map = n0 + np.zeros_like(cmam)
    else:
        n_map = n0 * exp(-cmam)
    return n_map

'''
def scattering_ma(event_type, p, maia_d, row, col):
    """Return the Rayleigh or Compton differential mass attenuation coefficients
    (cm2/g/sr) for icru44 brain tissue with density described by the 'matrix'
    map of phantom p into the Maia detector element indexed by row, col. This
    needs to be multiplied later by the Maia-channel-dependent solid angle.

    Arguments:
    event_type - string, one of ['rayleigh', 'compton']
    p - phantom instance (matrix plus elements)
    maia_d - Maia instance

    row, col - maia detector element indices

    Returns:
    The (rayleigh, compton) differential mass attenuation coefficient

    """
    assert event_type in ['rayleigh', 'compton']

    # Get spherical angles (polar theta & azimuthal phi) to detector element
    theta = maia_d.pads[q].theta
    phi = maia_d.pads[q].phi

    compound = p.matrix.cp  # elemental data for matrix compound

    # Get the contribution to Rayleigh and Compton scattering from each
    # element in the matrix compound and sum these.
    ma = 0.0
    for el in compound:
        # Assuming propagation along z, coordinate system used by the DCSP_Rayl
        # and DCSP_Compt methods is shown here:
        # http://upload.wikimedia.org/wikipedia/commons/thumb/4/4f/
        # 3D_Spherical.svg/200px-3D_Spherical.svg.png
        # i.e. spherical coordinates with polar angle theta, azimuthal angle phi
        z = xrl.SymbolToAtomicNumber(compound[el])

        # Mass attenuation coefficients from cross-sections. See
        # http://physics.nist.gov/PhysRefData/XrayMassCoef/chap2.html
        # Units of the following expression:
        # elemental fraction by weight * differential mass attenuation coefft
        # unitless * cm2/g/sr
        if event_type == 'rayleigh':
            f = xrl.DCSP_Rayl
        else:
            f = xrl.DCSP_Compt
        ma += compound[el] * f(z, p.energy, theta, phi)
    return ma

def compton_scattered_energy(energy_in, q, maia_d):
    """Energy of the Compton photons scattered into the direction of the Maia
    detector element.

    Parameters
    ----------
    energy_in : float
        Incident beam photon energy (keV).
    q : int
        Maia detector channel id
    maia_d : Maia instance

    Returns
    -------
    float
        Energy of scattered photons (keV)

    """
    # Get polar scattering angle (theta) to detector element
    # Get spherical angles (polar theta & azimuthal phi) to detector element
    theta = maia_d.pads[q].theta

    # energy_out = 1.0 / (1.0/energy_in +
    # (1 - np.cos(theta))*J_PER_KEV/sc.m_e/sc.c/sc.c)
    return xrl.ComptonEnergy(energy_in, theta)
'''

'''
def old_emission_map(event_type, p, n_map, angle, el=None):
    """Compute the maia-detector-shaped map of Rayleigh, Compton or K-edge
    fluorescence from the element map for an incident irradiance map n_map.

    Parameters
    ----------
    event_type : string
        One of ['rayleigh', 'compton', 'fluoro'].
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    n_map : 2d ndarray of float
        Map of incident irradiance.
    angle : float
        Stage rotation angle (degrees).
    el : string, optional
        Name of element (e.g. 'Fe') used if projecting that element's
        fluorescence.

    Returns
    -------
    1d ndarray of float
        The accumulated flux in Maia detector row 7 for the requested edge.

    """
    assert event_type in ['rayleigh', 'compton', 'fluoro']

    # Get matrix map and, for fluorescence, the map for the requested element,
    # and rotate them to the same angle as the intensity map since these all
    # need to be in registration.
    matrix_map = zero_outside_circle(p.el_maps['matrix'])
    matrix_map_r = rotate(matrix_map, -angle)
    del matrix_map

    k_alpha_energy = -1
    if event_type == 'fluoro':
        edge_map = zero_outside_circle(p.el_maps[el])
        edge_map_r = rotate(edge_map, -angle)
        del edge_map

        # Do this here because we're outside the detector channel loop.
        # Get Z for the fluorescing element and check that its K_alpha is
        # below the incident energy.
        el_z = xrl.SymbolToAtomicNumber(el)
        line = xrl.KA_LINE
        k_alpha_energy = xrl.LineEnergy(el_z, line)
        assert k_alpha_energy < p.energy

    # 2d accumulator for results
    accumulator = np.empty((maia_d.shape[1], n_map.shape[0]))

    # Iterate over maia detector elements in theta, i.e. maia columns
    # This should be parallelizable
    row = 7
    # Get multiplication factor for additional distance that an outgoing photon
    # must travel as it passes out-of-plane to the detector.
    delta_theta_yx_radian = maia_d.yx_angles_radian(row, col=0)
    delta_theta_y_radian = delta_theta_yx_radian[0]
    y_distance_factor = 1.0 / np.cos(delta_theta_y_radian)

    for channel_id in maia_d.channel_selection(row=row):
        col = maia_d.maia_data_column_from_id(channel_id, 'Column')
        # Get angle to rotate maps so that propagation toward detector plane
        # corresponds with direction to maia detector element
        delta_theta_yx_radian = maia_d.yx_angles_radian(row, col)
        delta_theta_x = rad_to_deg(delta_theta_yx_radian[1])

        # For every maia detector element, get the solid angle (parallelize?)
        # Orient the maps toward the maia element
        # TODO: check sign of delta: +ve or -ve?
        imap_rm = rotate(n_map, -delta_theta_x)
        matrix_map_rm = rotate(matrix_map_r, -delta_theta_x)

        # Rotate the geometry so that the detector is on the bottom, so we can
        # integrate by stepping through the row indices.
        imap_rm = np.rot90(imap_rm)
        matrix_map_rm = np.rot90(matrix_map_rm)

        # Solid angle of Maia channel
        omega = maia_d.solid_angle(row, col)

        # mass attenuation coefft. (cm2/g)
        if event_type == 'fluoro':
            edge_map_rm = rotate(edge_map_r, -delta_theta_x)
            edge_map_rm = np.rot90(edge_map_rm)

            # Simulate fluorescence event:
            # Generate the initial fluorescence intensity
            # Start with the highest-energy edge or a mean factor accounting for
            # all edges first (move to other edges later?)
            mac = fluoro_ma(p, el_z)

            # Scale for propagation over one voxel
            # *_mac_t = *_mac * p.um_per_px/UM_PER_CM
            # (cm3/g/sr) =   (cm2/g/sr) * cm
            mac_t = mac * p.um_per_px / UM_PER_CM * y_distance_factor
            # Generate outgoing radiation.
            # This is the fluorescence intensity map.
            imap_rm *= -expm1(-edge_map_rm * omega * mac_t)
            del edge_map_rm
        else:
            mac = scattering_ma(event_type, p, row, col)
            mac_t = mac * p.um_per_px / UM_PER_CM * y_distance_factor
            # Generate outgoing radiation.
            # This is the scattering radiation intensity map.
            imap_rm *= -expm1(-matrix_map_rm * omega * mac_t)

        # Now we've "evented," we use the mass attenuation coefficients of the
        # matrix for propagation with absorption out to the detector, but this
        # is at a new energy depending on the event type.
        energy = {
            'rayleigh': p.energy,
            'compton': compton_scattered_energy(p.energy, row, col),
            'fluoro': k_alpha_energy,
        }[event_type]
        mac_t = brain.ma(energy) * p.um_per_px / UM_PER_CM * y_distance_factor

        # Propagate all intensity to the detector, accumulating (+) and
        # absorbing [exp(-mu/rho rho t)] as we go.
        cmam_matrix = np.cumsum(matrix_map_rm, axis=0) * mac_t
        i_out = (imap_rm * exp(-cmam_matrix * omega)).sum(axis=0)

        # Store the result for this detector element.
        accumulator[col] = i_out

    return accumulator
'''

def outgoing_photon_energy(event_type, p, q=None, maia_d=None, el=None):
    """Return the interaction-type-dependent outgoing photon energy in keV.

    Parameters
    ----------
    event_type : string
        One of ['rayleigh', 'compton', 'fluoro'].
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    q : int (optional, required for compton interaction)
        Maia detector channel id
    maia_d : Maia() instance (optional, required for compton interaction)
    el : string (optional, required for fluoro interaction)
        Name of fluorescing element (e.g. 'Fe').

    Returns
    -------
    float
        Outgoing photon energy (keV)

    """
    assert event_type in ['rayleigh', 'compton', 'fluoro']

    def k_alpha_energy(el):
        el_z = xrl.SymbolToAtomicNumber(el)
        line = xrl.KA_LINE
        energy = xrl.LineEnergy(el_z, line)
        # assert energy < p.energy
        return energy

    '''
    if event_type == 'rayleigh':
        energy = p.energy
    elif event_type == 'compton':
        energy = compton_scattered_energy(p.energy, q, maia_d)
    else:           # 'fluoro'
    '''
    if True:        # Just support 'fluoro' for now
        energy = k_alpha_energy(el)

    return energy

def fluoro_emission_map(p, n_map, angle, el):
    """Compute the maia-detector-shaped map of K-edge
    fluorescence from the element map for an incident irradiance map n_map.

    Parameters
    ----------
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    n_map : 2d ndarray of float
        Map of incident irradiance.
    angle : float
        Stage rotation angle (degrees).
    el : string
        Name of fluorescing element (e.g. 'Fe').

    Returns
    -------
    2d ndarray of float
        The fluorescence emission map for the requested edge.

    """
    # edge_map = zero_outside_circle(p.el_maps[el])
    edge_map = p.el_maps[el]
    edge_map_r = rotate(edge_map, angle)
    del edge_map

    # Get Z for the fluorescing element
    el_z = xrl.SymbolToAtomicNumber(el)
    line = xrl.KA_LINE

    # Sanity check that el K_alpha is below the incident energy.
    k_alpha_energy = xrl.LineEnergy(el_z, line)
    # assert k_alpha_energy < p.energy
    if k_alpha_energy >= p.energy:
        Q = 0.0
    else:
        # Simulate fluorescence event:
        # CS_FluorLine_Kissel_Cascade is the XRF cross section Q_{i,YX} in Eq. (12)
        # of Schoonjans et al.
        Q = xrl.CS_FluorLine_Kissel_Cascade(el_z, line, p.energy)
        # print(el, end=' ')

    # 2d array for results
    emission_map = n_map * Q * edge_map_r * p.um_per_px / UM_PER_CM

    return emission_map

def channel_emission_map(event_type, p, q, maia_d, n_map, angle, el):
    """Select and defer to the interaction-type-specific emission map
    computation.

    Parameters
    ----------
    event_type : string
        One of ['rayleigh', 'compton', 'fluoro'].
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    q : int
        Maia detector channel id
    maia_d : Maia() instance
    n_map : 2d ndarray of float
        Map of incident irradiance.
    angle : float
        Stage rotation angle (degrees).
    el : string
        Name of fluorescing element (e.g. 'Fe').

    Returns
    -------
    2d ndarray of float
        The fluorescence emission map for the requested edge.

    """
    assert event_type in ['rayleigh', 'compton', 'fluoro']

    if event_type == 'fluoro':
        e_map = channel_fluoro_map(p, q, maia_d, n_map, angle, el)
    elif event_type == 'rayleigh':
        e_map = channel_rayleigh_map(p, q, maia_d, n_map, angle)
    elif event_type == 'compton':
        e_map = channel_compton_map(p, q, maia_d, n_map, angle)
    else:
        raise RuntimeError('Brain explodes')

    return e_map

def channel_fluoro_map(p, q, maia_d, n_map, angle, el):
    """Compute the maia-detector-shaped map of K-edge
    fluorescence from the element map for an incident irradiance map n_map.

    Parameters
    ----------
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    q : int
        Maia detector channel id
    maia_d : Maia instance
    n_map : 2d ndarray of float
        Map of incident irradiance.
    angle : float
        Stage rotation angle (degrees).
    el : string
        Name of fluorescing element (e.g. 'Fe').

    Returns
    -------
    2d ndarray of float
        The fluorescence emission map for the requested edge.

    """
    solid_angle = maia_d.pads[q].omega
    return (solid_angle / 4 / np.pi *
            fluoro_emission_map(p, n_map, angle, el))

def channel_rayleigh_map(p, q, maia_d, n_map, angle):
    """Compute the maia-detector-shaped map of Rayleigh scattering
    from the element map for an incident irradiance map n_map.

    Parameters
    ----------
    p : Phantom2d object
        p.energy - incident beam photon energy (keV).
        p.um_per_px - length of one pixel of the map (um).
    q : int
        Maia detector channel id
    maia_d : Maia() instance
    n_map : 2d ndarray of float
        Map of incident irradiance.
    angle : float
        Stage rotation angle (degrees).

    Returns
    -------
    2d ndarray of float
        The fluorescence emission map for the requested edge.

    """
    solid_angle = maia_d.pads[q].omega

    # Get spherical angles (polar theta & azimuthal phi) to detector element
    theta = maia_d.pads[q].theta
    phi = maia_d.pads[q].phi

    energy = outgoing_photon_energy('rayleigh', p)

    # get a list of all elements el and their weights w_el

    # The absorption map mu = ma_M * mu_M + sum_k ( ma_k * mu_k )
    mu = p.matrix.ma(energy) * p.el_maps['matrix']
    for el in p.el_maps:
        if el == 'matrix':
            continue
        Z = xrl.SymbolToAtomicNumber(el)
        mu += p.matrix.cp[el] * xrl.DCSP_Rayl(Z, energy, theta, phi) * \
              p.el_maps[el]

    ma += xrl.DCSP_Rayl(Z, energy, theta, phi)

    # edge_map = zero_outside_circle(p.el_maps[el])
    edge_map_r = rotate(edge_map, angle)
    del edge_map

    # Get Z for the fluorescing element
    el_z = xrl.SymbolToAtomicNumber(el)
    line = xrl.KA_LINE

    # Sanity check that el K_alpha is below the incident energy.
    k_alpha_energy = xrl.LineEnergy(el_z, line)
    assert k_alpha_energy < p.energy

    # Simulate fluorescence event:
    # CS_FluorLine_Kissel_Cascade is the XRF cross section Q_{i,YX} in Eq. (12)
    # of Schoonjans et al.
    Q = xrl.CS_FluorLine_Kissel_Cascade(el_z, line, p.energy)

    # 2d array for results
    emission_map = n_map * Q * edge_map_r * p.um_per_px / UM_PER_CM

    return emission_map

def write_sinogram(im, p, event_type, el='matrix'):
    """Project and write sinogram for element map el
    Creates a filename by prepending s_ to the filename and appending r or c
    if writing the rayleigh or compton sinogram, respectively.

    Parameters
    ----------
    im : 2d ndarray of float
        sinogram.
    p : Phantom2d object
    event_type : string
        One of ['absorption', 'rayleigh', 'compton', 'fluoro'].
    el : string
        name of current element, e.g. 'Fe'. (default 'matrix').

    """
    # Get the filename that matches the glob pattern for this element
    # and prepend s_ to it
    pattern = p.filename

    matches = helpers.match_pattern(pattern, glob.glob(pattern))
    if matches:
        # Just get the filename that's a match for the element el
        match_base = [m[0] for m in matches if el==m[1]][0]
    else:
        raise Exception('Element {} not found in {}'.format(el, matches))
    path = os.path.dirname(pattern)
    base = os.path.basename(match_base)
    s_filename = os.path.join(path, 's_'+base)

    # Write sinogram (absorption map)

    # append r, c, f, a to -matrix suffix so sinograms read in as unique images
    if '-matrix' in s_filename:
        s_filename = s_filename.replace('-matrix', '-matrix-' + event_type[0])
    write_tiff32(s_filename, im)

    return im

import numpy,from numpy,import scipy,from scipy


# hst.py

import os
import numpy as np
import pandas as pd

from ..io.read_hst import read_hst
from ..load_sim import LoadSim

class Hst:

    @LoadSim.Decorators.check_pickle_hst
    def read_hst(self, savdir=None, merge_mhd=True, force_override=False):
        """Function to read hst and convert quantities to convenient units
        """

        hst = read_hst(self.files['hst'], force_override=force_override)

        # delete the first row (post-processing)
        hst.drop(hst.index[:1], inplace=True)

        u = self.u
        domain = self.domain

        # volume of resolution element (code unit)
        dvol = domain['dx'].prod()
        # total volume of domain (code unit)
        vol = domain['Lx'].prod()
        # Area of domain (code unit)
        LxLy = domain['Lx'][0]*domain['Lx'][1]

        # Time in code unit
        hst['time_code'] = hst['time']
        # Time in Myr
        hst['time'] *= u.Myr
        # Total gas mass in Msun
        hst['mass'] *= vol*u.Msun
        # Gas surface density in Msun/pc^2
        hst['Sigma_gas'] = hst['mass']/(LxLy*u.pc**2)
        # H mass/surface density in Msun
        hst['MH'] = hst['mass']/u.muH
        hst['Sigma_H'] = hst['MH']/(LxLy*u.pc**2)

        # H2 mass in Msun
        try:
            hst['MH2'] = hst['MH2']*vol*u.Msun/u.muH
            hst['Sigma_H2'] = hst['MH2']/(LxLy*u.pc**2)
            hst['MH2_noLW'] = hst['MH2_noLW']*vol*u.Msun/u.muH
            hst['Sigma_H2_noLW'] = hst['MH2_noLW']/(LxLy*u.pc**2)
        except KeyError:
            pass

        # Neutral gas mass in Msun
        try:
            hst['Mneu'] = hst['scalar{:d}'.format(domain['IHI'])]*vol*u.Msun
        except KeyError:
            pass

        # Star formation rate per area [Msun/kpc^2/yr]
        hst['sfr10'] = hst['sfr10']/(LxLy*u.pc**2)
        hst['sfr40'] = hst['sfr40']/(LxLy*u.pc**2)
        hst['sfr100'] = hst['sfr100']/(LxLy*u.pc**2)

        # Cosmic ray ionization rate without attenuation
        hst['xi_CR0'] = 2e-16*self.par['problem']['xi_CR_amp']*(hst['sfr40']/3e-3)

        hst.index = hst['time_code']
        #hst.index.name = 'index'

        # Merge with mhd history dump
        if merge_mhd:
            try:
                hst_mhd = self.read_hst_mhd()
                hst = hst_mhd.reindex(hst.index, method='nearest',
                                      tolerance=0.1).combine_first(hst)
            except FileNotFoundError:
                pass

        self.hst = hst

        return hst

        # # Ionized gas mass in Msun
        # hst['Mion'] *= vol*u.Msun
        # # Collisionally ionized gas (before ray tracing) in Msun
        # hst['Mion_coll'] *= vol*u.Msun
        # # Total photoionization rate [#/sec]
        # hst['Qiphot'] *= vol*(u.length**3).cgs
        # # Total collisional ionization rate [#/sec]
        # hst['Qicoll'] *= vol*(u.length**3).cgs
        # # Total dust absorption rate [#/sec]
        # hst['Qidust'] *= vol*(u.length**3).cgs

    #     # Mass fraction ionized gas
    #     hst['mf_ion'] = hst['Mion']/hst['mass']
    #     hst['mf_ion_coll'] = hst['Mion_coll']/hst['mass']

        # for f in range(self.par['radps']['nfreq']):
        #     # Total luminosity [Lsun]
        #     hst['Ltot_cl{:d}'.format(f)] *= vol*u.Lsun
        #     hst['Ltot_ru{:d}'.format(f)] *= vol*u.Lsun
        #     hst['Ltot{:d}'.format(f)] = \
        #         hst['Ltot_cl{:d}'.format(f)] + hst['Ltot_ru{:d}'.format(f)]
        #     # Total luminosity included in simulation
        #     hst['L_cl{:d}'.format(f)] *= vol*u.Lsun
        #     hst['L_ru{:d}'.format(f)] *= vol*u.Lsun
        #     hst['L{:d}'.format(f)] = \
        #         hst['L_cl{:d}'.format(f)] + hst['L_ru{:d}'.format(f)]
        #     # Luminosity that escaped boundary
        #     hst['Lesc{:d}'.format(f)] *= vol*u.Lsun
        #     # Luminosity lost due to dmax
        #     hst['Llost{:d}'.format(f)] *= vol*u.Lsun
        #     # Escape fraction, lost fraction
        #     # Estimation of true escape fraction estimation (upper bound)
        #     hst['fesc{:d}'.format(f)] = hst['Lesc{:d}'.format(f)] / \
        #                                 hst['L{:d}'.format(f)]
        #     hst['flost{:d}'.format(f)] = hst['Llost{:d}'.format(f)] / \
        #                                  hst['L{:d}'.format(f)]
        #     hst['fesc{:d}_est'.format(f)] = hst['fesc{:d}'.format(f)] + \
        #                                     hst['flost{:d}'.format(f)]
        #     hst['fesc{:d}_cum_est'.format(f)] = \
        #         (hst['Lesc{:d}'.format(f)] + hst['Llost{:d}'.format(f)]).cumsum() / \
        #          hst['L{:d}'.format(f)].cumsum()

        # return hst

    #     # Scale heights of [warm] ionized gas, nesq
    #     # Check if columns exist

    #     # nesq
    #     if 'H2nesq' in hst.columns and 'nesq' in hst.columns:
    #         hst['H_nesq'] = np.sqrt(hst['H2nesq'] / hst['nesq'])
    #         hst.drop(columns=['H2nesq', 'nesq'], inplace=True)

    #     # Warm nesq
    #     if 'H2wnesq' in hst.columns and 'wnesq' in hst.columns:
    #         hst['H_wnesq'] = np.sqrt(hst['H2wnesq'] / hst['wnesq'])
    #         hst.drop(columns=['H2wnesq', 'wnesq'], inplace=True)

    #     # For warm medium,
    #     # append _ to distinguish from mhd history variable
    #     if 'H2w' in hst.columns and 'massw' in hst.columns:
    #         hst['H_w_'] = np.sqrt(hst['H2w'] / hst['massw'])
    #         hst['Mw_'] = hst['massw']*vol*u.Msun
    #         hst['mf_w_'] = hst['Mw_']/hst['mass']
    #         hst.drop(columns=['H2w', 'massw'], inplace=True)

    #     # Warm ionized
    #     if 'H2wi' in hst.columns and 'masswi' in hst.columns:
    #         hst['H_wi'] = np.sqrt(hst['H2wi'] / hst['masswi'])
    #         hst['Mwion'] = hst['masswi']*vol*u.Msun
    #         hst['mf_wion'] = hst['Mwion']/hst['mass']
    #         hst.drop(columns=['H2wi', 'masswi'], inplace=True)

    #     ##########################
    #     # With ionizing radiation
    #     ##########################
    #     if self.par['radps']['nfreq'] == 2 and \
    #        self.par['radps']['nfreq_ion'] == 1:
    #         hnu0 = self.par['radps']['hnu[0]']/u.eV
    #         hnu1 = self.par['radps']['hnu[1]']/u.eV
    #         # Total luminosity
    #         hst['Qitot_cl'] = hst['Ltot_cl0']/u.Lsun/hnu0/u.s
    #         hst['Qitot_ru'] = hst['Ltot_ru0']/u.Lsun/hnu0/u.s
    #         hst['Qitot'] = hst['Qitot_ru'] + hst['Qitot_cl']
    #         # Total Q included as source
    #         hst['Qi_cl'] = hst['L_cl0']/u.Lsun/hnu0/u.s
    #         hst['Qi_ru'] = hst['L_ru0']/u.Lsun/hnu0/u.s
    #         hst['Qi'] = hst['Qi_ru'] + hst['Qi_cl']
    #         hst['Qiesc'] = hst['Lesc0']/u.Lsun/hnu0/u.s
    #         hst['Qilost'] = hst['Llost0']/u.Lsun/hnu0/u.s
    #         hst['Qiesc_est'] = hst['Qilost'] + hst['Qiesc']

    #     else:
    #         self.logger.error('Unrecognized option nfreq={0:d}, nfreq_ion={1:d}'.\
    #                           format(self.par['radps']['nfreq'],
    #                                  self.par['radps']['nfreq_ion']))

    #     # midplane radiation energy density in cgs units
    #     hst['Erad0_mid'] *= u.energy_density
    #     hst['Erad1_mid'] *= u.energy_density

    #     try:
    #         hst.to_pickle(fpkl)
    #     except IOError:
    #         self.logger.warning('[read_hst]: Could not pickle hst to {0:s}.'.format(fpkl))

    #     self.hst = hst
    #     return self.hst

    def read_hst_mhd(self):

        # Read original mhd history dump from /tigress/changgoo
        hst = read_hst('/tigress/changgoo/{0:s}/hst/{0:s}.hst'.\
                       format(self.problem_id), force_override=True)

        u = self.u
        domain = self.par['domain1']
        Lx = domain['x1max'] - domain['x1min']
        Ly = domain['x2max'] - domain['x2min']
        Lz = domain['x3max'] - domain['x3min']
        Nx = domain['Nx1']
        Ny = domain['Nx2']
        Nz = domain['Nx3']
        Ntot = Nx*Ny*Nz
        vol = Lx*Ly*Lz
        LxLy = Lx*Ly
        dz = Lz/Nz
        Omega = self.par['problem']['Omega']
        time_orb = 2*np.pi/Omega*u.Myr # Orbital time in Myr

        if 'x1Me' in hst:
            mhd = True
        else:
            mhd = False

        h = pd.DataFrame()
        h['time_code'] = hst['time']
        h['time'] = h['time_code']*u.Myr # time in Myr
        h['time_orb'] = h['time']/time_orb

        h['mass'] = hst['mass']*u.Msun*vol
        h['Sigma'] = h['mass']/LxLy
        h['mass_sp'] = hst['msp']*u.Msun*vol
        h['Sigma_sp'] = h['mass_sp']/LxLy

        # Mass, volume fraction, scale height
        h['H'] = np.sqrt(hst['H2'] / hst['mass'])
        for ph in ['c','u','w','h1','h2']:
            h['mf_{}'.format(ph)] = hst['M{}'.format(ph)]/hst['mass']
            h['vf_{}'.format(ph)] = hst['V{}'.format(ph)]
            h['H_{}'.format(ph)] = \
                np.sqrt(hst['H2{}'.format(ph)] / hst['M{}'.format(ph)])

        # mf, vf, H of thermally bistable (cold + unstable + warm) medium
        h['mf_2p'] = h['mf_c'] + h['mf_u'] + h['mf_w']
        h['vf_2p'] = h['vf_c'] + h['vf_u'] + h['vf_w']
        h['H_2p'] = np.sqrt((hst['H2c'] + hst['H2u'] + hst['H2w']) / \
                            (hst['Mc'] + hst['Mu'] + hst['Mw']))

        # Kinetic and magnetic energy
        h['KE'] = hst['x1KE'] + hst['x2KE'] + hst['x3KE']
        if mhd:
            h['ME'] = hst['x1ME'] + hst['x2ME'] + hst['x3ME']

        hst['x2KE'] = hst['x2dke']
        for ax in ('1','2','3'):
            Ekf = 'x{}KE'.format(ax)
            if ax == '2':
                Ekf = 'x2dke'
            # Mass weighted velocity dispersion??
            h['v{}'.format(ax)] = np.sqrt(2*hst[Ekf]/hst['mass'])
            if mhd:
                h['vA{}'.format(ax)] = \
                    np.sqrt(2*hst['x{}ME'.format(ax)]/hst['mass'])
            h['v{}_2p'.format(ax)] = \
                np.sqrt(2*hst['x{}KE_2p'.format(ax)]/hst['mass']/h['mf_2p'])

        h['cs'] = np.sqrt(hst['P']/hst['mass'])
        h['Pth_mid'] = hst['Pth']*u.pok
        h['Pth_mid_2p'] = hst['Pth_2p']*u.pok/hst['Vmid_2p']
        h['Pturb_mid'] = hst['Pturb']*u.pok
        h['Pturb_mid_2p'] = hst['Pturb_2p']*u.pok/hst['Vmid_2p']

        # Midplane number density
        h['nmid'] = hst['nmid']
        h['nmid_2p'] = hst['nmid_2p']/hst['Vmid_2p']

        # Star formation rate per unit area [Msun/kpc^2/yr]
        h['sfr10']=hst['sfr10']
        h['sfr40']=hst['sfr40']
        h['sfr100']=hst['sfr100']

        h.index = h['time_code']
        #h.index.name = 'index'

        self.hst_mhd = h

        return self.hst_mhd

        # return pd.read_pickle(
        #     '/tigress/changgoo/{0:s}/hst/{0:s}.hst_cal.p'.format(self.problem_id))


#!/usr/bin/env python

"""
    Does some simple calculations to test trace gas  PACE calculations

    Adapted from benchmark4Amir.py
    Patricia Castellanos, April 2020

"""

import os
import sys
from   netCDF4 import Dataset
from   netCDF4 import Dataset as ncread
import numpy   as np
from MAPL.constants import *
from py_leo_vlidort import VLIDORT_POLAR_
from scipy.interpolate import interp1d
import scipy.integrate as integrate
from pyhdf.SD import SD, SDC
from multiprocessing import Pool

format   = 'NETCDF4_CLASSIC'
plane_parallel = True
MISSING = -1.e+20

WrapperFuncs = {'MODIS_BRDF'                : VLIDORT_POLAR_.vector_brdf_modis,
                'MODIS_BRDF_BPDF'           : VLIDORT_POLAR_.vector_brdf_modis_bpdf,
                'BPDF'                      : VLIDORT_POLAR_.vector_bpdf,
                'LAMBERTIAN'                : VLIDORT_POLAR_.vector_lambert,
                'LAMBERTIAN_BPDF'           : VLIDORT_POLAR_.vector_lambert_bpdf,
                'GissCX'                    : VLIDORT_POLAR_.vector_gisscx,
                'CX'                        : VLIDORT_POLAR_.vector_cx,
                'OCICX'                     : VLIDORT_POLAR_.vector_ocicx,
                'OCIGissCX'                 : VLIDORT_POLAR_.vector_ocigisscx,
                'OCIGissCX_NOBM_CLOUD'      : VLIDORT_POLAR_.vector_ocigisscx_nobm_cloud,
                'ROT_CALC'                  : VLIDORT_POLAR_.rot_calc}

def get_geom(Iscan,Icross):
    inFile = '/nobackup/PACE/L1B/Y2020/M03/D24/OCI2020084005000.L1B_PACE.nc'
    nc = Dataset(inFile)
    grp = nc.groups['geolocation_data']
    sza = grp.variables['solar_zenith'][:]
    saa = grp.variables['solar_azimuth'][:]
    vza = grp.variables['sensor_zenith'][:]
    vaa = grp.variables['sensor_azimuth'][:]

    # make azimuths clockwise from north
    I = saa < 0
    saa[I] = 360. + saa[I]

    I = vaa < 0
    vaa[I] = 360. + vaa[I]

    # define SAA according to photon travel direction
    saa = saa + 180.0
    I = saa >= 360.
    saa[I] = saa[I] - 360.

    raa = vaa - saa

    I = raa < 0
    raa[I] = raa[I] + 360.0

    saa = grp.variables['solar_azimuth'][:]
    vaa = grp.variables['sensor_azimuth'][:]

    sza = np.array([sza[Iscan,Icross]])
    vza = np.array([vza[Iscan,Icross]])
    raa = np.array([raa[Iscan,Icross]])
    saa = np.array([saa[Iscan,Icross]])
    vaa = np.array([vaa[Iscan,Icross]])

    return sza,vza,raa,saa,vaa

def get_ROT(ch,pe,te,ze,verbose=True):
    """
    calculate ROT and depol ratio
    """

    args = [ch, pe, ze, te, MISSING, verbose]
    ROT, depol_ratio, rc = VLIDORT_POLAR_.rot_calc(*args)

    return ROT,depol_ratio

def get_TOA_unpack(args):
    return get_TOA(*args)

def get_TOA(channel,
            F0,
            ROT,depol_ratio,
            tau,ssa,pmom,
            alpha,
            SZA,VZA,RAA,
            km,pe,te,ze,
            nstreams,
            albedoType,
            sleave,
            U10m,V10m,mr,
            verbose):
    """
    Do RT calculation to get TOA
    """

    # wrapper function based on albedo
    vlidortWrapper = WrapperFuncs[albedoType]

    if albedoType == 'OCIGissCX':
        args = [channel, nstreams, plane_parallel, ROT, depol_ratio, alpha, tau, ssa, pmom,
                pe, ze, te,
                U10m, V10m, mr,
                SZA, RAA, VZA,
                F0,
                MISSING,
                verbose]
        I, reflectance, surf_reflectance, Q, U, BR_Q, BR_U, rc = vlidortWrapper(*args)
    else:
        args = [channel, nstreams, plane_parallel, ROT, depol_ratio, alpha, tau, ssa, pmom,
                tau, ssa, pmom,
                tau, ssa, pmom,
                pe, ze, te,
                U10m, V10m, mr,
                sleave, False,
                SZA, RAA, VZA,
                F0,
                MISSING,
                verbose]
        I, reflectance, surf_reflectance, Q, U, BR_Q, BR_U, rc, adjusted_sleave = vlidortWrapper(*args)

    return I,reflectance,surf_reflectance

def writenc(outFile,channels,
            F0,
            SZA,SAA,VZA,VAA,
            I,BR,reflectance,
            ROD,depol_ratio,
            alphaD,
            pe,te,ze,
            U10m,V10m,mr):

    nch = len(channels)
    km  = len(pe)

    # Write data to netcdf file
    nc = Dataset(outFile,'w',format='NETCDF4_CLASSIC')
    nc.title = 'Line-by-Line TOA radiance calculation for one OCI pixel. Model 5 - US Standard 1962'

    dc = nc.createDimension('channels',nch)
    dn = nc.createDimension('npixel',1)
    dk = nc.createDimension('leve',km)

    ch = nc.createVariable('channels','f4',('channels',))
    ch.long_name = "Wavelength in nm"
    ch[:] = channels

    f  = nc.createVariable('solar_irradiance','f4',('channels',))
    f.long_name = 'Thuillier solar irradiance spectrum'
    f.units = 'uW/cm^2/nm'
    f[:] = F0

    rad = nc.createVariable('I','f4',('channels',))
    rad.long_name = "sun normalized TOA radiance"
    rad[:] = I

    ref = nc.createVariable('R','f4',('channels',))
    ref.long_name = "TOA reflectance"
    ref[:] = reflectance

    sr = nc.createVariable('surface_reflectance','f4',('channels',))
    sr.long_name = "surface bidirectional reflectance"
    sr[:] = BR

    a = nc.createVariable('sza','f4',('npixel',))
    a.long_name = "solar zenith angle"
    a.units = 'degrees'
    a[:] = SZA

    a = nc.createVariable('saa','f4',('npixel',))
    a.long_name = "solar azimiuth angle"
    a.units = 'degrees 0-360 clockwise from north'
    a[:] = SAA

    a = nc.createVariable('vza','f4',('npixel',))
    a.long_name = "sensor zenith angle"
    a.units = 'degrees'
    a[:] = VZA

    a = nc.createVariable('vaa','f4',('npixel',))
    a.long_name = "sensor azimiuth angle"
    a.units = 'degrees 0-360 clockwise from north'
    a[:] = SAA

    rod = nc.createVariable('ROD','f4',('channels',))
    rod.long_name = "Rayleigh Optical Depth"
    rod[:] = ROD

    depol = nc.createVariable('depol_ratio','f4',('channels',))
    depol.long_name = 'Rayleigh depolarization ratio'
    depol[:] = depol_ratio

    alpha = nc.createVariable('ALPHA','f4',('channels',))
    alpha.long_name = "Trace Gas Absorption Optical Depth"
    alpha[:] = alphaD

    ke = nc.createVariable('PE','f4',('leve',))
    ke.long_name = 'Pressure at layer edge'
    ke.units     = 'Pa'
    ke[:] = pe

    ke = nc.createVariable('TE','f4',('leve',))
    ke.long_name = 'Temperature at layer edge'
    ke.units     = 'K'
    ke[:] = te

    ke = nc.createVariable('ZE','f4',('leve',))
    ke.long_name = 'height above surface'
    ke.units = 'm'
    ke[:] = ze

    wi = nc.createVariable('U10M','f4',('npixel',))
    wi.long_name = 'U10M wind speed'
    wi.units = 'm/s'
    wi[:] = U10m

    wi = nc.createVariable('V10M','f4',('npixel',))
    wi.long_name = 'V10M wind speed'
    wi.units = 'm/s'
    wi[:] = V10m

    wi = nc.createVariable('mr','f4',('npixel',))
    wi.long_name = 'ocean water refractive index'
    wi[:] = mr

    nc.close()

def get_PTWV_profile(inFile,model=5):
    """
    Read in height [km],pressure [mb], temperature [K], water vapor vmr profile [ppm]
    for selected model:
    0:Tropical
    1:Mid Latitude Summer
    2:Mid Latitude Winter
    3:Subarctic Summer
    4:Subarctic Winter
    5:US Standard 1962
    6:User Defined Model
    """
    nc = Dataset(inFile)
    # make array because interp doesn't take masked arrays
    pe = np.array(nc.variables['p'][model,:])
    te = np.array(nc.variables['t'][model,:])
    ze = np.array(nc.variables['h'][model,:])
    vmre =np.array( nc.variables['vmr'][model,:])

    km   = len(ze) - 1

    # get DP, from Amir's code
    DP= np.empty([km])
    for i in range(km):
        DP[i] = (pe[i] - pe[i+1])/1013

    # convert mb to Pascal
    pe = pe *100
    DELP = pe[:-1] - pe[1:]

    # get T middles
    DELP = pe[:-1] - pe[1:]
    pm = pe[:-1] - 0.5*DELP
    f = interp1d(pe,te)
    tm = f(pm)

    # convert km to m
    ze = ze*1000.
    dz = ze[1:] - ze[:-1]

    # calculate air number density [molecules/m3]
    g = 9.80616 # gravity m/s2
    Na = 6.022e23 # avogadro's number
    MW_AIR =  28.964*1e-3 # kg/mole

    AIRDENS = DELP/(dz*g) # kg/m3
    rho = AIRDENS/MW_AIR      # moles/m3
    rho = rho*Na  #[molecules/m3]

    return km, pe, te, tm, ze, dz, rho, vmre, DP, AIRDENS

def get_abs(inFile):
    """
    Read aborption coefficients from Amir's HITRAN calculations
    """
    nc = Dataset(inFile)
    # wavenumber [cm-1]
    waveno = np.array(nc.variables['waveno'][:])
    # abosrption coefficients [not sure about units]
    abs_o2  = nc.variables['abscf_o2'][:]
    abs_h2o = nc.variables['abscf_h2o'][:]
    abs_co  = nc.variables['abscf_co'][:]
    abs_co2 = nc.variables['abscf_co2'][:]
    abs_ch4 = nc.variables['abscf_ch4'][:]
    abs_n2o = nc.variables['abscf_n2o'][:]

    nc.close()

    return abs_o2, abs_h2o, abs_co, abs_co2, abs_ch4, abs_n2o, waveno

def get_rsr(inFile):
    """
    Read in OCI RSR File
    """
    hdf = SD(inFile, SDC.READ)
    rsr = hdf.select('RSR')[:]
    wav_rsr = hdf.select('rsrwave')[:]
    wav_oci = hdf.select('wave')[:]
    hdf.end()

    return rsr, wav_rsr, wav_oci

def get_alpha(A,VMR,rho,dz):
    """
    Calculate Absorption optical depth profile
    A - absorption coefficient [m2/molecule]
    VMR - trace gas mixing ratio [vol/vol, dimensionless]
    rho - air number density [molecules/m3]
    ze  - profile layer thickness [m]
    """

    # convert vmr to molecules/m3
    nxe = VMR*rho

    # get the optical depth subcolumns
    km, nch = A.shape
    alpha = np.zeros([km,nch])
    for i in range(km):
        alpha[i,:] = nxe[i]*dz[i]*A[i,:]

    return alpha

# ---
def read_o3(inFile,te):
    f = open(inFile)
    nhead = 10
    for i in range(nhead):
        hh = f.readline()

    wav_o3, c0, c1, c2 = [],[],[],[]
    for l in f:
        a = np.array(l.split()).astype(float)
        wav_o3.append(a[0])
        c0.append(a[1])
        c1.append(a[2])
        c2.append(a[3])
    f.close()

    wav_o3 = np.array(wav_o3)
    c0     = np.array(c0)
    c1     = np.array(c1)
    c2     = np.array(c2)
    T0     = 273.15

    # calculate xsec for te
    nwav = len(wav_o3)
    km   = len(te)
    xsec_o3 = np.zeros([km,nwav])

    for i,t in enumerate(te):
        xsec_o3[i,:] = c0 + c1*(t-T0) + c2*(t-T0)**2

    xsec_o3 = xsec_o3*1e-20

    # convert from cm2/molecule to m2/molecule
    xsec_o3 = xsec_o3*1e-4

    return wav_o3, xsec_o3

# --
def read_ROD_table(inFile):
    f = open(inFile)

    for i in range(16):
        f.readline() #header

    wav = []
    rod = []
    depol = []
    for l in f:
        w, r, d = l.split()
        wav.append(w)
        rod.append(r)
        depol.append(d)

    f.close()

    wav = np.array(wav).astype('float')
    rod = np.array(rod).astype('float')
    depol = np.array(depol).astype('float')

    return wav, rod, depol

#------------------------------------ M A I N ------------------------------------

if __name__ == "__main__":

    outRoot = 'hyperTest/'
    outFile   = '{}/outputs/hyperTest_CK_Thuillier_g5nr.nc4'.format(outRoot)

    # Pressure [Pa], temperature [K], height [m], water vapor [ppm]  profile - standard atmosphere
    # used to make OCI look up tables
    # rhoe = air number density [molecules/m3]
    inFile    = '{}/atrem_tpvmr.nc'.format(outRoot)
    km, pe, te, tm, ze, dz, rho, h2oe, DP, AIRDENS =  get_PTWV_profile(inFile)

    # Read in G5NR CO, CO2, O3, RH
    Iscan  = 600
    Icross = 1000
    LevelB = '/nobackup/PACE/LevelB/Y2006/M03/D24'
    inFile = '{}/pace-g5nr-std.lb.chm_Nv.20060324_005000.nc4'.format(LevelB)
    nc = Dataset(inFile)
    co = nc.variables['CO'][0,:,Iscan,Icross]    # VMR
    co2 = nc.variables['CO2'][0,:,Iscan,Icross]  # VMR
    o3 = nc.variables['O3'][0,:,Iscan,Icross]  #kg/kg
    nc.close()
    inFile = '{}/pace-g5nr-std.lb.aer_Nv.20060324_005000.nc4'.format(LevelB)
    nc = Dataset(inFile)
    h2o = nc.variables['WV_VMR'][0,:,Iscan,Icross]  #ppm
    nc.close()

    # flip from bottom to top
    co  = co[::-1]
    co2 = co2[::-1]
    o3  = o3[::-1]
    h2o = h2o[::-1]

    # Read in alpha_table
    inFile = 'alphaTable_v0/alpha_CK_Thuillier_o3.nc4'
    nc = Dataset(inFile)
    all_wl = np.array(nc.variables['channels'][:])
    g_bins   = nc.variables['g_bins'][:]
    alpha_o2 = nc.variables['alpha_o2'][:]
    ROD      = nc.variables['ROD'][:]
    depol    = nc.variables['depol_ratio'][:]
    F0_int   = nc.variables['solar_irradiance'][:]

    nc.close()

    alpha = alpha_o2

#    # integrate air density in each layer
#    rhoint = rho*dz

    # get absorption optical depth with new aborption coefficient
#    co_vmr = 0.1*1e-6
#    co_vmr = co
#    alpha_co = get_alpha(abs_co_z,co_vmr,rho,dz)

#    o2_vmr = 0.21
#    alpha_o2 = get_alpha(abs_o2_z,o2_vmr,rho,dz)

#    co2_vmr = 400.*1.0E-06
#    co2_vmr = co2
#    alpha_co2 = get_alpha(abs_co2_z,co2_vmr,rho,dz)

#    ch4_vmr = 1.8*1.0E-06
#    alpha_ch4 = get_alpha(abs_ch4_z,ch4_vmr,rho,dz)

#    n2o_vmr = 0.3*1.0E-06
#    alpha_n2o = get_alpha(abs_n2o_z,n2o_vmr,rho,dz)

#    h2o_vmr = h2o*1.0e-6
#    alpha_h2o = get_alpha(abs_h2o_z,h2o_vmr,rho,dz)

    # add up all the alphas
#    alpha = alpha_h2o + alpha_n2o + alpha_ch4 + alpha_co2 + alpha_o2 + alpha_co

#    # ----
#    # OZONE Stuff
#    # ---
#    # read xsec
#    inFile = 'hyperTest/o3_bremen/Ozone_abs_x_wTemperatureFit.dat'
#    # wav [nm], C0, C1(T), C2(T^2)
#    # xsec is in m2/molecule
#    wav_o3,abs_o3 = read_o3(inFile,tm)
#    # reverse so going from max to min wavelength
#    wav_o3 = wav_o3[::-1]
#    abs_o3 = abs_o3[:,::-1]

#    # interpolate to LBL wavelengths
#    # append zeros to max lbl wavelength
#    wav_new = np.arange(wav_o3.max()+0.1,wav_abs.max()+0.1,0.1)
#    # reverse
#    wav_new = wav_new[::-1]
#    wav_o3  = np.append(wav_new,wav_o3)
#    nnew = len(wav_new)
#    abs_o3 = np.append(np.zeros([km,nnew]),abs_o3,axis=1)

#    abs_o3_lbl = np.zeros(abs_h2o_z.shape)
#    for k in range(km):
#       # xsec_f = interp1d(wav_o3,abs_o3[k,:],kind='linear')
#        abs_o3_lbl[k,:] = xsec_f(wav_abs)

#    # append UV-Vis to LBL that stops at 555
#    i = wav_o3 < wav_abs.min()
#    all_wl = np.append(wav_abs,wav_o3[i])
#    abs_o3_lbl = np.append(abs_o3_lbl,abs_o3[:,i],axis=1)

#    # convert mass mixing ratio to molecules/m3
#    Na = 6.022e23 # avogadro's number
#    O3_MW  = 48.0*1e-3     # kg/mole

#    o3_conc = o3*AIRDENS   # kg/m3
#    o3_conc = o3_conc*Na/O3_MW   # molecules/m3

#    # get the optical depth subcolumns
#    alpha_o3 = np.zeros(abs_o3_lbl.shape)
#    for i in range(km):
#        alpha_o3[i,:] = o3_conc[i]*dz[i]*abs_o3_lbl[i,:]

    nwav = len(all_wl)

#    # add ozone to total alpha
#    # extend array down to uv
#    nnew = nwav - len(wav_abs)
#    alpha = np.append(alpha,np.zeros([km,nnew]),axis=1)
#    alpha = alpha + alpha_o3

#    # limit to wavelengths covered by RSR
#    I = all_wl <= wav_rsr.max()+1.0
#    all_wl = all_wl[I]
#    alpha   = alpha[:,I]
#    alpha[:] = 0.0

#    # get ROD from oci_tables
#    inFile = 'oci_tables/rayleigh_bodhaine.txt'
#    wav,rod,depol = read_ROD_table(inFile)

#    # interpolate to lbl wavelengths
#    rod_f = interp1d(wav,rod,kind='linear')
#    rod_lbl = rod_f(all_wl)
#    depol_f = interp1d(wav,depol,kind='linear')
#    depol_lbl = depol_f(all_wl)

    # flip everything vertically so going from top of atmosphere to surface
    pe = pe[-1::-1]
    te = te[-1::-1]
    ze = ze[-1::-1]
#    alpha = alpha[-1::-1,:]

    # add dimension to be in km+1,nobs
    pe.shape = (km+1,1)
    te.shape = (km+1,1)
    ze.shape = (km+1,1)

    # read in granule geometry
    SZA,VZA,RAA,SAA,VAA = get_geom(Iscan,Icross)
    csza = np.cos(np.radians(SZA))

#    # Read in solar irradiance spectrum
#    # second dim = wavelength, irradiance
#    # units=nm, uW/cm^2/nm
#    inFile = '{}/Thuillier_F0.npy'.format(outRoot)
#    F0 = np.load(inFile)
#    # interpolate to wavelengths
#    F0_f = interp1d(F0[:,0],F0[:,1],kind='linear',fill_value="extrapolate")
#    F0_int = F0_f(all_wl)

    # --------------
    # surface stuff
    # --------------

    # SLEAVE
    LevelB = '/nobackup/PACE/LevelB/surface/SLEAVE/NOBM/Y2006/M03/D24'
    inFile = '{}/pace-g5nr.lb.sleave.20060324_005000.nc4'.format(LevelB)
    nc     = Dataset(inFile)
    nobm_wav = nc.variables['wavelength'][:]
    rrs      = nc.variables['rrs'][0,:,Iscan,Icross]
    nc.close()
    rrs_f    = interp1d(nobm_wav,rrs,kind='linear',fill_value=0.0,bounds_error=False)
    sleave   = rrs_f(all_wl)*csza

    # wind speed
    LevelB = '/nobackup/PACE/LevelB/Y2006/M03/D24'
    inFile = '{}/pace-g5nr-std.lb.met_Nv.20060324_005000.nc4'.format(LevelB)
    nc     = Dataset(inFile)
    U10m   = np.array([nc.variables['U10M'][0,Iscan,Icross]])
    V10m   = np.array([nc.variables['V10M'][0,Iscan,Icross]])

    # water refractive index
    mr   = 1.334

    # RT stuff
    nstreams = 12
    albedoType = 'OCIGissCX_NOBM_CLOUD'
    albedoType = 'OCIGissCX'

    # loop through channels
    nproc = 50
    nwl   = len(all_wl)
#    nwl   = 100
#    wlstep = int(nwl/nproc)
    wlstep  = 10
    sys.exit()
    args = []
    ROD  = []
    depol = []
    sys.exit()
    for ich in np.arange(0,nwl,wlstep):
        endch = ich + wlstep
        if endch > nwl:
            endch = nwl
        nch = endch - ich
        ch = all_wl[ich:endch]

        # Get Rayleigh
        # ROT shape is nlev,1,nch
        ROT, depol_ratio = get_ROT(ch,pe,te,ze,verbose=False)

        ROT = ROT*rod_lbl[ich:endch]/np.squeeze(ROT.sum(axis=0))
        ROD.append(np.squeeze(ROT.sum(axis=0)))

        depol_ratio = depol_lbl[ich:endch]
        depol.append(depol_ratio)

        # trace gas
        alpha_ch = alpha[:,ich:endch]
        alpha_ch.shape = (km,1,nch)

        # AOP vectors [km,nch,nobs]
        tau = np.zeros([km,nch,1])
        ssa = np.zeros([km,nch,1])
        pmom = np.zeros([km,nch,1,30,6])

        # water refractive index
        mr_in = np.ones(nch)
        mr_in = mr_in*mr

        # solar irradiance
        F0_in = F0_int[ich:endch]
        F0_in.shape = (nch,1)

        # sleave
        sleave_ch = sleave[ich:endch]

        args.append([ch,
                     F0_in,
                     ROT,depol_ratio,
                     tau,ssa,pmom,
                     alpha_ch,
                     SZA,VZA,RAA,
                     km,pe,te,ze,
                     nstreams,
                     albedoType,
                     sleave_ch,
                     U10m,V10m,mr_in,
                     False])

#       I,reflectance,BR =  get_TOA(ch,
#                                    ROT,depol_ratio,
#                                    tau,ssa,pmom,
#                                    alpha_ch,
#                                    SZA,VZA,RAA,
#                                    km,pe,te,ze,
#                                    nstreams,
#                                    albedoType,
#                                    U10m=U10m,V10m=V10m,mr=np.array([1.334]),
#                                    verbose=False)

    # use multiprocessing
    p = Pool(nproc)
    result = p.map(get_TOA_unpack,args)
    I = []
    reflectance = []
    BR = []
    for r in result:
        I_r,reflectance_r,BR_r = r
        I.append(np.squeeze(I_r))
        reflectance.append(np.squeeze(reflectance_r))
        BR.append(np.squeeze(BR_r))

    p.close()
    p.join()

    # concatenate arrays
    ROD = np.concatenate(ROD)
    depol_ratio = np.concatenate(depol)

    I  = np.concatenate(I)
    reflectance = np.concatenate(reflectance)
    BR = np.concatenate(BR)

    alphaD = alpha.sum(axis=0)
    # write to outFile
    writenc(outFile,all_wl,
            F0_int,
            SZA,SAA,VZA,VAA,
            I,BR,reflectance,
            ROD,depol_ratio,
            alphaD,
            pe,te,ze,
            U10m,V10m,mr)


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

################################################################################
#
#   CanTherm
#
#   Copyright (c) 2010 by Joshua W. Allen (jwallen@mit.edu)
#
#   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.
#
################################################################################

import math
import numpy
import logging

from rmgpy.quantity import constants
from rmgpy.statmech import *

################################################################################

def applyEnergyCorrections(E0, modelChemistry, atoms, bonds):
    """
    Given an energy `E0` in J/mol as read from the output of a quantum chemistry
    calculation at a given `modelChemistry`, adjust the energy such that it
    is consistent with the normal gas-phase reference states. `atoms` is a
    dictionary associating element symbols with the number of that element in
    the molecule. `bonds` is a dictionary associating bond types with the number
    of that bond in the molecule.
    """

    # Spin orbit correction (SOC) in Hartrees
    # Values taken from note 22 of http://jcp.aip.org/resource/1/jcpsa6/v109/i24/p10570_s1 and converted to hartrees
    # Values in millihartree are also available (with fewer significant figures) from http://jcp.aip.org/resource/1/jcpsa6/v106/i3/p1063_s1
    SOC = {'H':0.0, 'N':0.0, 'O': -0.000355, 'C': -0.000135, 'P': 0.0, 'S': -0.000893}

    # Step 1: Reference all energies to a model chemistry-independent basis
    # by subtracting out that model chemistry's atomic energies
    # Note: If your model chemistry does not include spin orbit coupling, you should add the corrections to the energies here
    if modelChemistry == 'CBS-QB3':
        atomEnergies = {'H':-0.499818 , 'N':-54.520543, 'O':-74.987624, 'C':-37.785385, 'P':-340.817186, 'S': -397.657360}
    elif modelChemistry == 'G3':
        atomEnergies = {'H':-0.5010030, 'N':-54.564343, 'O':-75.030991, 'C':-37.827717, 'P':-341.116432}
    elif modelChemistry == 'Klip_1':
        atomEnergies = {'H':-0.50003976 + SOC['H'], 'O':-75.00915718 + SOC['O'], 'C':-37.79249556 + SOC['C']}
    elif modelChemistry == 'Klip_2':
        #Klip QCI(tz,qz)
        atomEnergies = {'H':-0.50003976 + SOC['H'], 'O':-75.00692746 + SOC['O'], 'C':-37.79044863 + SOC['C']}
    elif modelChemistry == 'Klip_2_cc':
        #Klip CCSD(T)(tz,qz)
        atomEnergies = {'H':-0.50003976 + SOC['H'], 'O':-75.00681155 + SOC['O'], 'C':-37.79029443 + SOC['C']}
    else:
        logging.warning('Unknown model chemistry "{0}"; not applying energy corrections.'.format(modelChemistry))
        return E0
    for symbol, count in atoms.iteritems():
        if symbol in atomEnergies: E0 -= count * atomEnergies[symbol] * 4.35974394e-18 * constants.Na
        else:
            logging.warning('Ignored unknown atom type "{0}".'.format(symbol))

    # Step 2: Atom energy corrections to reach gas-phase reference state
    # Experimental enthalpy of formation at 0 K
    # See Gaussian thermo whitepaper at http://www.gaussian.com/g_whitepap/thermo.htm)
    # Note: these values are relatively old and some improvement may be possible by using newer values, particularly for carbon
    # However, care should be taken to ensure that they are compatible with the BAC values (if BACs are used)
    atomHf = {'H': 51.63 , 'N': 112.53 ,'O': 58.99 ,'C': 169.98, 'S': 65.66 }
    # Thermal contribution to enthalpy Hss(298 K) - Hss(0 K) reported by Gaussian thermo whitepaper
    # This will be subtracted from the corresponding value in atomHf to produce an enthalpy used in calculating the enthalpy of formation at 298 K
    atomThermal = {'H': 1.01 , 'N': 1.04, 'O': 1.04 ,'C': 0.25, 'S': 1.05 }
    # Total energy correction used to reach gas-phase reference state
    # Note: Spin orbit coupling no longer included in these energies, since some model chemistries include it automatically
    atomEnergies = {}
    for element in atomHf:
        atomEnergies[element] = atomHf[element] - atomThermal[element]
    for symbol, count in atoms.iteritems():
        if symbol in atomEnergies: E0 += count * atomEnergies[symbol] * 4184

    # Step 3: Bond energy corrections
    bondEnergies = { 'C-H': -0.11, 'C-C': -0.3, 'C=C': -0.08, 'C#C': -0.64,
        'O-H': 0.02, 'C-O': 0.33, 'C=O': 0.55, 'N#N': -2.0, 'O=O': -0.2,
        'H-H': 1.1, 'C#N': -0.89, 'S-H': 0.0, 'C-S': 0.43, 'S=O': -0.78 }
    for symbol, count in bonds.iteritems():
        if symbol in bondEnergies: E0 += count * bondEnergies[symbol] * 4184
        else:
            logging.warning('Ignored unknown bond type "{0}".'.format(symbol))

    return E0

################################################################################

def projectRotors(geom, F, rotors, linear, TS):
    """
    For a given geometry `geom` with associated force constant matrix `F`,
    lists of rotor information `rotors`, `pivots`, and `top1`, and the linearity
    of the molecule `linear`, project out the nonvibrational modes from the
    force constant matrix and use this to determine the vibrational frequencies.
    The list of vibrational frequencies is returned in cm^-1.
    """

    Nrotors = len(rotors)
    Natoms = len(geom.mass)
    Nvib = 3 * Natoms - (5 if linear else 6) - Nrotors - (1 if (TS) else 0)

    if linear:
        D = numpy.zeros((Natoms*3,5+Nrotors), numpy.float64)
    else:
        D = numpy.zeros((Natoms*3,6+Nrotors), numpy.float64)

    for i in range(Natoms):
        # Projection vectors for translation
        D[3*i+0,0] = 1.0
        D[3*i+1,1] = 1.0
        D[3*i+2,2] = 1.0
        # Projection vectors for [external] rotation
        D[3*i:3*i+3,3] = numpy.array([0, -geom.coordinates[i,2], geom.coordinates[i,1]], numpy.float64)
        D[3*i:3*i+3,4] = numpy.array([geom.coordinates[i,2], 0, -geom.coordinates[i,0]], numpy.float64)
        if not linear:
            D[3*i:3*i+3,5] = numpy.array([-geom.coordinates[i,1], geom.coordinates[i,0], 0], numpy.float64)
    for i, rotor in enumerate(rotors):
        scanLog, pivots, top, symmetry = rotor
        # Determine pivot atom
        if pivots[0] in top: pivot = pivots[0]
        elif pivots[1] in top: pivot = pivots[1]
        else: raise Exception('Could not determine pivot atom.')
        # Projection vectors for internal rotation
        e12 = geom.coordinates[pivots[0],:] - geom.coordinates[pivots[1],:]
        e12 /= numpy.linalg.norm(e12)
        for atom in top:
            e31 = geom.coordinates[atom,:] - geom.coordinates[pivot,:]
            D[3*atom:3*atom+3,-Nrotors+i] = numpy.cross(e31, e12)

    # Make sure projection matrix is orthonormal
    import scipy.linalg
    D = scipy.linalg.orth(D)

    # Project out the non-vibrational modes from the force constant matrix
    P = numpy.dot(D, D.transpose())
    I = numpy.identity(Natoms*3, numpy.float64)
    F = numpy.dot(I - P, numpy.dot(F, I - P))

    # Generate mass-weighted force constant matrix
    # This converts the axes to mass-weighted Cartesian axes
    # Units of Fm are J/m^2*kg = 1/s^2
    Fm = F.copy()
    for i in range(Natoms):
        for j in range(Natoms):
            for u in range(3):
                for v in range(3):
                    Fm[3*i+u,3*j+v] /= math.sqrt(geom.mass[i] * geom.mass[j]) / constants.Na

    # Get eigenvalues of mass-weighted force constant matrix
    eig, V = numpy.linalg.eigh(Fm)
    eig.sort()

    # Convert eigenvalues to vibrational frequencies in cm^-1
    # Only keep the modes that don't correspond to translation, rotation, or internal rotation
    return numpy.sqrt(eig[-Nvib:]) / (2 * math.pi * constants.c * 100)

################################################################################

def saveStates(species, geometry, label, path):
    """
    Append the molecular degrees of freedom for `species` with associated
    string `label` to the file located at `path` on disk.
    """

    f = open(path, 'a')

    coordinates = geometry.coordinates * 1e10
    number = geometry.number
    numbers = {1: 'H', 6: 'C', 7: 'N', 8: 'O', 14: 'Si', 15: 'P', 16: 'S'}

    f = open(path, 'a')
    f.write('# Coordinates for {0} (angstroms):\n'.format(label))
    for i in range(coordinates.shape[0]):
        x = coordinates[i,0] - coordinates[0,0]
        y = coordinates[i,1] - coordinates[0,1]
        z = coordinates[i,2] - coordinates[0,2]
        f.write('#   {0} {1:9.4f} {2:9.4f} {3:9.4f}\n'.format(numbers[number[i]], x, y, z))

    f.write('states(\n')
    f.write('    label = "{0}",\n'.format(label))
    f.write('    E0 = {0!r},\n'.format(species.E0))
    f.write('    modes = [\n')
    for mode in species.states.modes:
        f.write('        {0!r},\n'.format(mode))
    f.write('    ],\n')
    f.write('    spinMultiplicity = {0:d},\n'.format(species.states.spinMultiplicity))
    try:
        f.write('    frequency={0!r},\n'.format(species.frequency))
    except AttributeError: pass
    f.write('    short_comment = "",\n')
    f.write('    long_comment = \n')
    f.write('"""\n')
    f.write('\n')
    f.write('""",\n')
    f.write(')\n\n')

    f.close()


# -*- coding: utf-8 -*-
"""

Machine Learning applied to defeat a 30 years old DOS game...

Game mechanics
"Imbiss" v. 5.4 by T. Bauer for IBM PC, Public Domain

Custom Environments in OpenAI’s Gym

Pocs (2020), Beginner’s Guide to Custom Environments in OpenAI’s Gym
How to set up, verify, and use a custom environment in reinforcement learning training with Python
https://towardsdatascience.com/beginners-guide-to-custom-environments-in-openai-s-gym-989371673952
https://github.com/MatePocs/gym-basic/blob/main/gym_basic/envs/basic_env_2.py

@author: Seinheilig, 2021
"""

import gym
import numpy as np

class ImbissEnv(gym.Env):

        def __init__(self):
            self.action_space = gym.spaces.Box(np.array([0,0,0,0,0,0,0]),np.array([300,300,300,300,1300,1300,1300]),dtype=np.int16)  # price tags of all 7 elements
            self.observation_space = gym.spaces.Box(np.array([1,-9]),np.array([7,40])) # day of week / temperature

            #### ToDo: Observations better MultiDiscrete?  // requre different logic for TE...
            ### self.observation_space = gym.spaces.MultiDiscrete([ 1, 49 ]) # day of week / temperature

        def rnd_state(self):
            #self.state = np.array([np.random.randint(1,7),np.random.randint(0,49)-9],dtype=np.int16) # full state (all days)
            self.state = np.array([1,np.random.randint(0,49)-9],dtype=np.int16) # only weekday...

        def step(self, action):
            '''
            single step
            return
            '''
            V = action  # price tags (actions)
            T = self.state[0] # day of the week
            TE = self.state[1] # temperatur
            K = 0 # money before selling
            reward =  self.Customer_Simulation(V,T,TE,K,Version=3)
            self.rnd_state()
            done = True
            info = {}
            return self.state, reward, done, info

        def reset(self):
            self.rnd_state()
            return self.state

        def render(self, mode='human'):
            pass

        def Kaufen(self,WS,V,S,K):
            """ reduced game mechanics for buying (DE: "kaufen")
            """
            debug_ = False
            H = [10,10,10,20,200,35,55]  # lowest price possible
            H = [10,10,10,20,200,50,70]  # good price..
            if debug_:
              Waren = ['Schokoeis','Vanilleeis','Erdbeereis','Cola','Zigaretten','Bratwurst','Pommes']
              print('Kunde: Könnten Sie mir bitte',S,Waren[WS],'geben')

            return K + S*V[WS] - S*H[WS]  # return new net balance after selling (income-expense)

        def Customer_Simulation(self,V,T,TE,K,Version=3):
            """ Spielmechanik - Kunden Simulation

            Args:
              V: Verkaufspreise
              T: Wochentag in [1:7]
              TE: Temperatur
              K: Kontostand vor der Simulation
              Version: 1) "Imbiss-Bude" von F. Brall 1983 für Apple II
                       2) "Imbiss" von O. Schwald 1984 für Commodore C64
                       3) "Imbiss" von T. Bauer 1991 für PC

            Returns:
              K: Kontostand nach der Simulation

            """
            debug_ = False
            if Version <3:
                EK = 10  # Eis Kunden
                ZK = 10  # Zigaretten Kunden
                BK = 30  # Bratwurst Kunden
                if T == 6:  # Samstag
                   EK = 15
                   ZK = 13
                   BK = 40
                if T == 7:  # Sonntag
                   EK = 20
                   ZK = 18
                   BK = 40

                # Korrektur der Kunden als Funktion des Preise
                EK -= int(np.min(V[0:3])/10) # bugfix!!
                ZK -= int(V[4]/100)
                BK -= int(np.min(V[5:7])/20) # bugfix!!

                # Temperatur Korrektur
                EK += int(TE/2)
                BK -= int(TE/2)
            else:
                ##### "Imbiss" PC 1991
                # 5 Änderungen zum Orginal:
                # 1) nutze max V der Warengruppe für Korrektur
                # 2) ZK = [10,12,15], ZK Basis Wochentag/Sa/So
                # 3) EK += 10 / andere Temperaturabhängigkeit
                # 4) BK andere Temperaturabhängigkeit
                # 5) if ZK/EK/BK < 0 --> ZK/EK/BK = 0 bevor AK berechnet wird
                EK = 20  # Eis Kunden
                ZK = 10  # Zigaretten Kunden
                BK = 30  # Bratwurst Kunden
                if T == 6:  # Samstag
                   EK = 25
                   ZK = 12
                   BK = 40
                if T == 7:  # Sonntag
                   EK = 30
                   ZK = 15
                   BK = 40

                # Korrektur der Kunden als Funktion des Preise
                EK -= int(np.max(V[0:4])/10)  # bugfix!!
                ZK -= int(V[4]/100)
                BK -= int(np.max(V[5:7])/20)  # bugfix!!

                # Temperatur Korrektur
                EK += int(TE/4)
                BK -= int(TE/3)

            # zu hoher Preis in einer Warengruppe (neg Kundenwert) führt nicht(!) zu einer Reduktion der Gesamtzahl der Kunden
            if BK < 0:
                BK = 0
            if ZK < 0:
                ZK = 0
            if EK < 0:
                EK = 0

            AK = ZK+BK+EK
            if debug_:
              print("Customer_Simulation:",V,T,TE,K)
              print('Kunden gesamt:\t\t',AK)
              print('Eis Kunden:\t\t',EK)
              print('Zigaret. Kunden:\t',ZK)
              print('Bratwurst Kunden:\t',BK)
            if AK < 0:
                AK = 0

            if AK < 2:
                return K

            for Kunde in range(AK):
                #print('Kunde',Kunde)
                kauf = False
                E = np.random.randint(0,9)    # eine Zufallsvariabe für Anzahl Einheiten zu kaufen und Preise ignorieren
                if E == 2:
                    S = 2
                else:
                    S = 1

                while not(kauf):
                    # eine zufällige Ware wählen
                    Z = np.random.randint(0,7)

                    if Z < 4 and  EK == 0:  # keine Eiskunden mehr
                        #print('keine Eiskunden mehr')
                        kauf = False
                        continue

                    if Z == 4 and  ZK == 0:  # keine Zigarettenkunden mehr
                        continue

                    if Z > 4 and  BK == 0:  # keine Bratwurstkunden mehr
                        continue

                    if Z == 0: # Schokoeis
                        if E != 3:  # Preisprüfung überspringen? (in 10% der Fälle)
                           if V[0]-V[1] > 20:
                               K = self.Kaufen(1,V,S,K)
                               EK -= 1; kauf = True
                               break
                           if V[0]-V[2] > 20:
                               K = self.Kaufen(2,V,S,K)
                               EK -= 1; kauf = True
                               break
                        K = self.Kaufen(0,V,S,K)
                        EK -= 1; kauf = True

                    if Z == 1: # Vanilleeis
                        if E != 3:  # Preisprüfung überspringen? (in 10% der Fälle)
                           if V[1]-V[0] > 20:
                               K = self.Kaufen(0,V,S,K)
                               EK -= 1; kauf = True
                               break
                           if V[1]-V[2] > 20:
                               K = self.Kaufen(2,V,S,K)
                               EK -= 1; kauf = True
                               break
                        K = self.Kaufen(1,V,S,K)
                        EK -= 1; kauf = True

                    if Z == 2: # Erdbeereis
                        if E != 3:  # Preisprüfung überspringen? (in 10% der Fälle)
                           if V[2]-V[0] > 20:
                               K = K = self.Kaufen(0,V,S,K)
                               EK -= 1; kauf = True
                               break
                           if V[2]-V[1] > 20:
                               K = self.Kaufen(1,V,S,K)
                               EK -= 1; kauf = True
                               break
                        K = self.Kaufen(2,V,S,K)
                        EK -= 1; kauf = True

                    if Z == 3: # Cola    ## KEINE PREISPRÜFUNG COLA!!!
                        K = self.Kaufen(3,V,S,K)
                        EK -= 1; kauf = True

                    if Z == 4: # Zigarette   ## PREISPRÜFUNG indirekt über ZK
                        K = self.Kaufen(4,V,S,K)
                        ZK -= 1; kauf = True

                    if Z == 5: # Bratwurst
                        if E != 3:  # Preisprüfung überspringen? (in 10% der Fälle)
                           if V[5]-V[6] > 30:
                               K = self.Kaufen(6,V,S,K)
                               BK -= 1; kauf = True
                               break
                        K = self.Kaufen(5,V,S,K)
                        BK -= 1; kauf = True

                    if Z == 6: # Pommes
                        if E != 3:  # Preisprüfung überspringen? (in 10% der Fälle)
                           if V[6]-V[5] > 30:
                               K = self.Kaufen(5,V,S,K)
                               BK -= 1; kauf = True
                               break
                        K = self.Kaufen(6,V,S,K)
                        BK -= 1; kauf = True

            return K