#  Copyright (C) 2014-2021 Syntrogi Inc dba Intheon. All rights reserved.

from pathlib import Path
import numpy as np
import os
import warnings

def parse_headmodel_file(filename):
    from scipy.io.matlab import loadmat
    from scipy.sparse import csc_matrix
    filepath = Path(filename)
    if not filepath.exists():
        filepath = Path(os.path.expanduser('~')) / '.stream_viewer' / 'headmodel' / filepath.name
    HM = loadmat(filepath, squeeze_me=True)
    leadfield = HM['leadfield']['matrix'][()]
    sensors = dict()
    sensors['labels'] = HM['sensors']['labels'][()]
    sensors['coordinates'] = HM['sensors']['coordinates'][()]
    laplacian = csc_matrix(HM['laplacian']['matrix'][()])

    meshes = [{'name': _, 'vertices': HM['meshes']['vertices'][ix], 'faces': HM['meshes']['faces'][ix] - 1}
              for ix, _ in enumerate(HM['meshes']['name'])]

    if HM['atlases'].size == 1:
        atlases = [{_: HM['atlases'][_][()] for _ in HM['atlases'].dtype.names}]
        atlases[0]['labeling'] -= 1
    else:
        atlases = [{'name': _, 'labeling': HM['atlases']['labeling'][ix] - 1, 'labels': HM['atlases']['labels'][ix]}
                   for ix, _ in enumerate(HM['atlases']['name'])]

    for atl in atlases:
        if atl['name'] == 'Desikan–Killiany' and len(atl['labeling']) == 4495:
            _labels = atl['labels'].tolist()

            atl['labeling'][202] = _labels.index('lingual L.1')  # From 'lingual R.2'
            switch_lr = [248, 334, 343, 1415, 1503, 1568, 1607, 1657, 1820, 1848, 2520, 2605, 3145, 3189, 3429, 3465,
                         3499, 3531, 3639, 3689, 3788, 3796, 3870, 3914, 3957, 4029, 4109, 4173, 4333]
            for v_id in switch_lr:
                old_roi_spl = _labels[atl['labeling'][v_id]].split(' ')
                new_roi = old_roi_spl[0] + ' ' + ('L' if old_roi_spl[1][0] == 'R' else 'R') + old_roi_spl[1][1:]
                atl['labeling'][v_id] = _labels.index(new_roi)

            # Previously known mislabels
            atl['labeling'][1038] = _labels.index('superiortemporal L.3')  # From middletemporal L.4
            atl['labeling'][985] = _labels.index('superiortemporal L.3')  # From supramarginal L.10
            atl['labeling'][1161] = _labels.index('superiortemporal L.4')  # From middletemporal L.6
            atl['labeling'][1509] = _labels.index('superiortemporal L.6')  # From precentral L.9
            atl['labeling'][1361] = _labels.index('superiortemporal L.5')  # From precentral L.9

    coords = HM['meta']['coordinates'][()]
    meta = {'coordinates': {_: coords[_][()] for _ in coords.dtype.names},
            'system': 'MNI'}
    del HM
    return {'leadfield': leadfield, 'sensors': sensors, 'laplacian': laplacian,
            'meshes': meshes, 'atlases': atlases, 'meta': meta}

def internalize_coordinates(coords, unit, x, y, z):
    """Convert the given coordinates (Nx3 array) into a unified
    internal coordinate system.

    The internal system is as follows:
    * unit=meters
    * x=right
    * y=front
    * z=up

    Args:
        coords: Nx3 array of coordinates to convert
        unit: the unit of the coordinates, e.g., 'meters', 'millimeters',
          can also be 'guess'
        x: orientation of the X axis relative to the head, e.g., 'front'
        y: orientation of the Y axis relative to the head, e.g., 'left'
        z: orientation of the Z axis relative to the head, e.g., 'up'

    Returns:
        coords: the transformed coordinates
    """

    # unit conversion
    if unit.lower() == 'guess':
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            max_coord = np.nanmax(np.abs(coords[:, :2]))
        if not np.isnan(max_coord):
            if max_coord < 0.1:
                unit = 'meters'
            elif max_coord < 1:
                unit = 'centimeters'
            elif max_coord < 10:
                unit = 'millimeters'

    if unit.lower() in ['millimeters', 'mm']:
        coords /= 1000.0
    elif unit.lower() in ['centimeters', 'cm']:
        coords /= 100.0
    elif not unit.lower() in ['meters', 'guess']:
        raise RuntimeError("Unsupported unit: %s" % unit)

    # rotations, etc
    if x == 'front' and y == 'left' and z == 'up':
        # +X nose direction
        # rotate 90 degrees clockwise (looking down on head)
        coords = np.dot(coords, np.array([[0, 1, 0],
                                          [-1, 0, 0],
                                          [0, 0, 1]]))
    elif x == 'back' and y == 'right' and z == 'up':
        # -X nose direction
        coords = np.dot(coords, np.array([[0, -1, 0],
                                          [1, 0, 0],
                                          [0, 0, 1]]))
    elif not (x == 'right' and y == 'front' and z == 'up'):
        # not +Y nose direction
        raise RuntimeError("Unsupported coordinate system (%s,%s,%s)" % (x, y,
                                                                         z))
    return coords


"""
Processing of GINI formatted data found on NOAAPORT
"""
import struct
import math
import zlib
from datetime import timezone, datetime
import os

import pyproj
import numpy as np
from pyiem.util import LOG

DATADIR = os.sep.join([os.path.dirname(__file__), "../data"])
M_PI_2 = 1.57079632679489661923
M_PI = 3.14159265358979323846
RE_METERS = 6371200.0
ENTITIES = [
    "UNK",
    "UNK",
    "MISC",
    "JERS",
    "ERS",
    "POES",
    "COMP",
    "DMSP",
    "GMS",
    "METEOSAT",
    "GOES7",
    "GOES8",
    "GOES9",
    "GOES10",
    "GOES11",
    "GOES12",
    "GOES13",
    "GOES14",
    "GOES15",
]
LABELS = [
    "UNK",
    "UNK",
    "MISC",
    "JERS",
    "ERS",
    "POES",
    "COMP",
    "DMSP",
    "GMS",
    "METEOSAT",
    "GOES",
    "GOES",
    "GOES",
    "GOES",
    "GOES",
    "GOES",
    "GOES",
    "GOES",
    "GOES",
]
CHANNELS = [
    "",
    "VIS",
    "3.9",
    "WV",
    "IR",
    "12",
    "13.3",
    "1.3",
    "U8",
    "U9",
    "U10",
    "U11",
    "U12",
    "LI",
    "PW",
    "SKIN",
    "CAPE",
    "TSURF",
    "WINDEX",
]
for _u in range(22, 100):
    CHANNELS.append(f"U{_u}")
SECTORS = [
    "NHCOMP",
    "EAST",
    "WEST",
    "AK",
    "AKNAT",
    "HI",
    "HINAT",
    "PR",
    "PRNAT",
    "SUPER",
    "NHCOMP",
    "CCONUS",
    "EFLOAT",
    "WFLOAT",
    "CFLOAT",
    "PFLOAT",
]

AWIPS_GRID_GUESS = {
    "A": 207,
    "B": 203,
    "E": 211,
    "F": 0,
    "H": 208,
    "I": 204,
    "N": 0,
    "P": 210,
    "Q": 205,
    "W": 211,
}

AWIPS_GRID = {
    "TIGB": 203,
    "TIGE": 211,
    "TIGW": 211,
    "TIGH": 208,
    "TIGP": 210,
    "TIGA": 207,
    "TIGI": 204,
    "TIGQ": 205,
    "TICF": 201,
}

def uint24(data):
    """convert three byte data that represents an unsigned int"""
    u = int(struct.unpack(">B", data[0:1])[0]) << 16
    u += int(struct.unpack(">B", data[1:2])[0]) << 8
    u += int(struct.unpack(">B", data[2:3])[0])
    return u

def int24(data):
    """Convert to int."""
    u = int(struct.unpack(">B", data[0:1])[0] & 127) << 16
    u += int(struct.unpack(">B", data[1:2])[0]) << 8
    u += int(struct.unpack(">B", data[2:3])[0])
    if (struct.unpack(">B", data[0:1])[0] & 128) != 0:
        u *= -1
    return u

def get_ir_ramp():
    """Return a np 256x3 array of colors to use for IR"""
    fn = "%s/gini_ir_ramp.txt" % (DATADIR,)
    data = np.zeros((256, 3), np.uint8)
    for i, line in enumerate(open(fn)):
        tokens = line.split()
        data[i, :] = [int(tokens[0]), int(tokens[1]), int(tokens[2])]
    return data

class GINIZFile:
    """
    Deal with compressed GINI files, which are the standard on NOAAPORT
    """

    def __init__(self, fobj):
        """Create a GNIFile instance with a compressed file object

        Args:
          fobj (file): A fileobject
        """
        fobj.seek(0)
        # WMO HEADER
        self.wmo = (fobj.read(21)).strip().decode("utf-8")
        d = zlib.decompressobj()
        hdata = d.decompress(fobj.read())
        self.metadata = self.read_header(hdata[21:])
        self.init_projection()
        totsz = len(d.unused_data)
        # 5120 value chunks, so we need to be careful!
        sdata = b""
        chunk = b"x\xda"
        i = 0
        for part in d.unused_data.split(b"x\xda"):
            if part == b"" and i == 0:
                continue
            chunk += part
            try:
                sdata += zlib.decompress(chunk)
                i += 1
                totsz -= len(chunk)
                chunk = b"x\xda"
            except Exception:
                chunk += b"x\xda"
        if totsz != 0:
            LOG.info("Totalsize left: %s", totsz)

        self.data = np.reshape(
            np.fromstring(sdata, np.int8),
            (self.metadata["numlines"] + 1, self.metadata["linesize"]),
        )

    def __str__(self):
        """return a string representation"""
        text = "%s Line Size: %s Num Lines: %s" % (
            self.wmo,
            self.metadata["linesize"],
            self.metadata["numlines"],
        )
        return text

    def awips_grid(self):
        """
        Return the awips grid number based on the WMO header
        """
        try1 = AWIPS_GRID.get(self.wmo[:4], None)
        if try1:
            return try1
        return AWIPS_GRID_GUESS.get(self.wmo[3], None)

    def current_filename(self):
        """
        Return a filename for this product, we'll use the format
        {SOURCE}_{SECTOR}_{CHANNEL}_{VALID}.png
        """
        return "%s_%s_%s.png" % (
            LABELS[self.metadata["creating_entity"]],
            SECTORS[self.metadata["sector"]],
            CHANNELS[self.metadata["channel"]],
        )

    def get_bird(self):
        """
        Return a string label for this satellite
        """
        return ENTITIES[self.metadata["creating_entity"]]

    def get_sector(self):
        """Return the sector."""
        return SECTORS[self.metadata["sector"]]

    def get_channel(self):
        """Return the channel."""
        return CHANNELS[self.metadata["channel"]]

    def archive_filename(self):
        """
        Return a filename for this product, we'll use the format
        {SOURCE}_{SECTOR}_{CHANNEL}_{VALID}.png
        """
        return ("%s_%s_%s_%s.png") % (
            LABELS[self.metadata["creating_entity"]],
            SECTORS[self.metadata["sector"]],
            CHANNELS[self.metadata["channel"]],
            self.metadata["valid"].strftime("%Y%m%d%H%M"),
        )

    def init_llc(self):
        """
        Initialize Lambert Conic Comformal
        """
        self.metadata["proj"] = pyproj.Proj(
            proj="lcc",
            lat_0=self.metadata["latin"],
            lat_1=self.metadata["latin"],
            lat_2=self.metadata["latin"],
            lon_0=self.metadata["lov"],
            a=6371200.0,
            b=6371200.0,
        )

        # s = 1.0
        # if self.metadata['proj_center_flag'] != 0:
        #    s = -1.0
        psi = M_PI_2 - abs(math.radians(self.metadata["latin"]))
        cos_psi = math.cos(psi)
        # r_E = RE_METERS / cos_psi
        alpha = math.pow(math.tan(psi / 2.0), cos_psi) / math.sin(psi)

        x0, y0 = self.metadata["proj"](
            self.metadata["lon1"], self.metadata["lat1"]
        )
        self.metadata["x0"] = x0
        self.metadata["y0"] = y0
        # self.metadata['dx'] *= alpha
        # self.metadata['dy'] *= alpha
        self.metadata["y1"] = y0 + (self.metadata["dy"] * self.metadata["ny"])

        (self.metadata["lon_ul"], self.metadata["lat_ul"]) = self.metadata[
            "proj"
        ](self.metadata["x0"], self.metadata["y1"], inverse=True)
        LOG.info(
            (
                "lat1: %.5f y0: %5.f y1: %.5f lat_ul: %.3f "
                "lat_ur: %.3f lon_ur: %.3f alpha: %.5f dy: %.3f"
            ),
            self.metadata["lat1"],
            y0,
            self.metadata["y1"],
            self.metadata["lat_ul"],
            self.metadata["lat_ur"],
            self.metadata["lon_ur"],
            alpha,
            self.metadata["dy"],
        )

    def init_mercator(self):
        """
        Compute mercator projection stuff
        """
        self.metadata["proj"] = pyproj.Proj(
            proj="merc",
            lat_ts=self.metadata["latin"],
            x_0=0,
            y_0=0,
            a=6371200.0,
            b=6371200.0,
        )
        x0, y0 = self.metadata["proj"](
            self.metadata["lon1"], self.metadata["lat1"]
        )
        self.metadata["x0"] = x0
        self.metadata["y0"] = y0

        x1, y1 = self.metadata["proj"](
            self.metadata["lon2"], self.metadata["lat2"]
        )
        self.metadata["x1"] = x1
        self.metadata["y1"] = y1

        self.metadata["dx"] = (x1 - x0) / self.metadata["nx"]
        self.metadata["dy"] = (y1 - y0) / self.metadata["ny"]

        (self.metadata["lon_ul"], self.metadata["lat_ul"]) = self.metadata[
            "proj"
        ](self.metadata["x0"], self.metadata["y1"], inverse=True)

        LOG.info(
            (
                "latin: %.2f lat_ul: %.3f lon_ul: %.3f "
                "y0: %5.f y1: %.5f dx: %.3f dy: %.3f"
            ),
            self.metadata["latin"],
            self.metadata["lat_ul"],
            self.metadata["lon_ul"],
            y0,
            y1,
            self.metadata["dx"],
            self.metadata["dy"],
        )

    def init_stereo(self):
        """
        Compute Polar Stereographic
        """
        self.metadata["proj"] = pyproj.Proj(
            proj="stere",
            lat_ts=60,
            lat_0=90,
            lon_0=self.metadata["lov"],
            x_0=0,
            y_0=0,
            a=6371200.0,
            b=6371200.0,
        )
        # First point!
        x0, y0 = self.metadata["proj"](
            self.metadata["lon1"], self.metadata["lat1"]
        )
        self.metadata["x0"] = x0
        self.metadata["y0"] = y0

        self.metadata["y1"] = y0 + (self.metadata["dy"] * self.metadata["ny"])
        (self.metadata["lon_ul"], self.metadata["lat_ul"]) = self.metadata[
            "proj"
        ](x0, self.metadata["y1"], inverse=True)

        LOG.info(
            (
                "lon_ul: %.2f lat_ul: %.2f "
                "lon_ll: %.2f lat_ll: %.2f "
                " lov: %.2f latin: %.2f lat1: %.2f lat2: %.2f "
                "y0: %5.f y1: %.5f dx: %.3f dy: %.3f"
            ),
            self.metadata["lon_ul"],
            self.metadata["lat_ul"],
            self.metadata["lon1"],
            self.metadata["lat1"],
            self.metadata["lov"],
            self.metadata["latin"],
            self.metadata["lat1"],
            self.metadata["lat2"],
            y0,
            self.metadata["y1"],
            self.metadata["dx"],
            self.metadata["dy"],
        )

    def init_projection(self):
        """
        Setup Grid and projection details
        """
        if self.metadata["map_projection"] == 3:
            self.init_llc()
        elif self.metadata["map_projection"] == 1:
            self.init_mercator()
        elif self.metadata["map_projection"] == 5:
            self.init_stereo()
        else:
            LOG.info("Unknown Projection: %s", self.metadata["map_projection"])

    def read_header(self, hdata):
        """read the header!"""
        meta = {}
        meta["source"] = struct.unpack("> B", hdata[0:1])[0]
        meta["creating_entity"] = struct.unpack("> B", hdata[1:2])[0]
        meta["sector"] = struct.unpack("> B", hdata[2:3])[0]
        meta["channel"] = struct.unpack("> B", hdata[3:4])[0]

        meta["numlines"] = struct.unpack(">H", hdata[4:6])[0]
        meta["linesize"] = struct.unpack(">H", hdata[6:8])[0]

        yr = 1900 + struct.unpack("> B", hdata[8:9])[0]
        mo = struct.unpack("> B", hdata[9:10])[0]
        dy = struct.unpack("> B", hdata[10:11])[0]
        hh = struct.unpack("> B", hdata[11:12])[0]
        mi = struct.unpack("> B", hdata[12:13])[0]
        ss = struct.unpack("> B", hdata[13:14])[0]
        # hs = struct.unpack("> B", hdata[14:15] )[0]
        meta["valid"] = datetime(yr, mo, dy, hh, mi, ss).replace(
            tzinfo=timezone.utc
        )
        meta["map_projection"] = struct.unpack("> B", hdata[15:16])[0]
        meta["proj_center_flag"] = struct.unpack("> B", hdata[36:37])[0] >> 7
        meta["scan_mode"] = struct.unpack("> B", hdata[37:38])[0]

        meta["nx"] = struct.unpack(">H", hdata[16:18])[0]
        meta["ny"] = struct.unpack(">H", hdata[18:20])[0]
        meta["res"] = struct.unpack(">B", hdata[41:42])[0]
        # Is Calibration Info included?
        # http://www.nws.noaa.gov/noaaport/document/ICD%20CH5-2005-1.pdf
        # page24
        # Mercator
        if meta["map_projection"] == 1:
            meta["lat1"] = int24(hdata[20:23])
            meta["lon1"] = int24(hdata[23:26])
            meta["lov"] = 0
            meta["dx"] = struct.unpack(">H", hdata[33:35])[0]
            meta["dy"] = struct.unpack(">H", hdata[35:37])[0]
            meta["latin"] = int24(hdata[38:41])
            meta["lat2"] = int24(hdata[27:30])
            meta["lon2"] = int24(hdata[30:33])
            meta["lat_ur"] = int24(hdata[55:58])
            meta["lon_ur"] = int24(hdata[58:61])
        # lambert == 3, polar == 5
        else:
            meta["lat1"] = int24(hdata[20:23])
            meta["lon1"] = int24(hdata[23:26])
            meta["lov"] = int24(hdata[27:30])
            meta["dx"] = uint24(hdata[30:33])
            meta["dy"] = uint24(hdata[33:36])
            meta["latin"] = int24(hdata[38:41])
            meta["lat2"] = 0
            meta["lon2"] = 0
            meta["lat_ur"] = int24(hdata[55:58])
            meta["lon_ur"] = int24(hdata[58:61])

        meta["dx"] = meta["dx"] / 10.0
        meta["dy"] = meta["dy"] / 10.0
        meta["lat1"] = meta["lat1"] / 10000.0
        meta["lon1"] = meta["lon1"] / 10000.0
        meta["lov"] = meta["lov"] / 10000.0
        meta["latin"] = meta["latin"] / 10000.0
        meta["lat2"] = meta["lat2"] / 10000.0
        meta["lon2"] = meta["lon2"] / 10000.0
        meta["lat_ur"] = meta["lat_ur"] / 10000.0
        meta["lon_ur"] = meta["lon_ur"] / 10000.0

        return meta


''' Compared with model_v1 use deeper network for 3D box regression. use BN for all FC layers
compared with model_v1_deeper, added more layers for 3D regression. added dropout for segmentation
compared with model_v1_deeper_0913, added visu for mean iou3d, added balanced loss.
compared with model_v1_deeper_0914, use two stage center regression..
'''
import tensorflow as tf
import numpy as np
import math
import sys
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(BASE_DIR, '../utils'))
import tf_util
from roi_seg_box3d_dataset import NUM_HEADING_BIN, NUM_SIZE_CLUSTER, NUM_CLASS, compute_box3d_iou, class2type, type_mean_size
from model_v1_deeper_0914_corners_balanced import get_box3d_corners, get_box3d_corners_helper # TODO move elsewhere
mean_size_arr = np.zeros((NUM_SIZE_CLUSTER, 3))
for i in range(NUM_SIZE_CLUSTER):
    mean_size_arr[i,:] = type_mean_size[class2type[i]]

def placeholder_inputs(batch_size, num_point):
    pointclouds_pl = tf.placeholder(tf.float32,
                                     shape=(batch_size, num_point, 6))
    one_hot_vec_pl = tf.placeholder(tf.float32, shape=(batch_size, NUM_CLASS))
    labels_pl = tf.placeholder(tf.int32,
                                shape=(batch_size, num_point))
    centers_pl = tf.placeholder(tf.float32,
                                shape=(batch_size, 3))
    heading_class_label_pl = tf.placeholder(tf.int32, shape=(batch_size,))
    heading_residual_label_pl = tf.placeholder(tf.float32, shape=(batch_size,))
    size_class_label_pl = tf.placeholder(tf.int32, shape=(batch_size,))
    size_residual_label_pl = tf.placeholder(tf.float32, shape=(batch_size,3))
    return pointclouds_pl, one_hot_vec_pl, labels_pl, centers_pl, heading_class_label_pl, heading_residual_label_pl, size_class_label_pl, size_residual_label_pl

def get_model(point_cloud, one_hot_vec, is_training, bn_decay=None):
    """ Classification PointNet, input is BxNx4, onehotvec is Bx3, output BxNx2 """
    batch_size = point_cloud.get_shape()[0].value
    num_point = point_cloud.get_shape()[1].value
    end_points = {}

    input_image = tf.expand_dims(point_cloud, -1)

    net = tf_util.conv2d(input_image, 64, [1,6],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv1', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 64, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv2', bn_decay=bn_decay)
    point_feat = tf_util.conv2d(net, 64, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv3', bn_decay=bn_decay)
    net = tf_util.conv2d(point_feat, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv4', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 1024, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv5', bn_decay=bn_decay)
    global_feat = tf_util.max_pool2d(net, [num_point,1],
                                     padding='VALID', scope='maxpool')
    print global_feat

    global_feat = tf.concat([global_feat, tf.expand_dims(tf.expand_dims(one_hot_vec, 1), 1)], axis=3)
    print 'Global Feat: ', global_feat
    global_feat_expand = tf.tile(global_feat, [1, num_point, 1, 1])
    print point_feat, global_feat_expand
    concat_feat = tf.concat(axis=3, values=[point_feat, global_feat_expand])
    print concat_feat

    net = tf_util.conv2d(concat_feat, 512, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv6', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 256, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv7', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv8', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv9', bn_decay=bn_decay)
    net = tf_util.dropout(net, is_training, 'dp1', keep_prob=0.5)

    logits = tf_util.conv2d(net, 2, [1,1],
                         padding='VALID', stride=[1,1], activation_fn=None,
                         scope='conv10')
    logits = tf.squeeze(logits, [2]) # BxNxC
    print logits

    print '-----------'
    #net = tf.concat(axis=3, values=[net, tf.expand_dims(tf.slice(point_cloud, [0,0,0], [-1,-1,3]), 2)])
    mask = tf.slice(logits,[0,0,0],[-1,-1,1]) < tf.slice(logits,[0,0,1],[-1,-1,1])
    mask = tf.to_float(mask) # BxNx1
    mask_count = tf.tile(tf.reduce_sum(mask,axis=1,keep_dims=True), [1,1,3]) # Bx1x3
    print mask
    point_cloud_xyz = tf.slice(point_cloud, [0,0,0], [-1,-1,3]) # BxNx3

    # ---- Subtract points mean ----
    mask_xyz_mean = tf.reduce_sum(tf.tile(mask, [1,1,3])*point_cloud_xyz, axis=1, keep_dims=True) # Bx1x3
    mask_xyz_mean = mask_xyz_mean/tf.maximum(mask_count,1) # Bx1x3
    point_cloud_xyz_stage1 = point_cloud_xyz - tf.tile(mask_xyz_mean, [1,num_point,1])
    print 'Point cloud xyz stage1: ', point_cloud_xyz_stage1

    # ---- Regress 1st stage center ----
    net = tf.expand_dims(point_cloud_xyz_stage1, 2)
    print net
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg1-stage1', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg2-stage1', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 256, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg3-stage1', bn_decay=bn_decay)
    mask_expand = tf.tile(tf.expand_dims(mask,-1), [1,1,1,256])
    masked_net = net*mask_expand
    print masked_net
    net = tf_util.max_pool2d(masked_net, [num_point,1], padding='VALID', scope='maxpool-stage1')
    net = tf.squeeze(net, axis=[1,2])
    print net
    net = tf.concat([net, one_hot_vec], axis=1)
    net = tf_util.fully_connected(net, 256, scope='fc1-stage1', bn=True, is_training=is_training, bn_decay=bn_decay)
    net = tf_util.fully_connected(net, 128, scope='fc2-stage1', bn=True, is_training=is_training, bn_decay=bn_decay)
    stage1_center = tf_util.fully_connected(net, 3, activation_fn=None, scope='fc3-stage1')
    stage1_center = stage1_center + tf.squeeze(mask_xyz_mean, axis=1) # Bx3
    end_points['stage1_center'] = stage1_center

    # ---- Subtract stage1 center ----
    point_cloud_xyz_submean = point_cloud_xyz - tf.expand_dims(stage1_center, 1)
    print 'Point cloud xyz submean: ', point_cloud_xyz_submean

    net = tf.expand_dims(point_cloud_xyz_submean, 2)
    print net
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg1', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg2', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 256, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg3', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 512, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv-reg4', bn_decay=bn_decay)
    mask_expand = tf.tile(tf.expand_dims(mask,-1), [1,1,1,512])
    masked_net = net*mask_expand
    print masked_net
    net = tf_util.max_pool2d(masked_net, [num_point,1], padding='VALID', scope='maxpool2')
    net = tf.squeeze(net, axis=[1,2])
    print net
    net = tf.concat([net, one_hot_vec], axis=1)
    net = tf_util.fully_connected(net, 512, scope='fc1', bn=True, is_training=is_training, bn_decay=bn_decay)
    net = tf_util.fully_connected(net, 256, scope='fc2', bn=True, is_training=is_training, bn_decay=bn_decay)

    # First 3 are cx,cy,cz, next NUM_HEADING_BIN*2 are for heading
    # next NUM_SIZE_CLUSTER*4 are for dimension
    output = tf_util.fully_connected(net, 3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4, activation_fn=None, scope='fc3')
    print output

    center = tf.slice(output, [0,0], [-1,3])
    center = center + stage1_center # Bx3
    end_points['center'] = center

    heading_scores = tf.slice(output, [0,3], [-1,NUM_HEADING_BIN])
    heading_residuals_normalized = tf.slice(output, [0,3+NUM_HEADING_BIN], [-1,NUM_HEADING_BIN])
    end_points['heading_scores'] = heading_scores # BxNUM_HEADING_BIN
    end_points['heading_residuals_normalized'] = heading_residuals_normalized # BxNUM_HEADING_BIN (should be -1 to 1)
    end_points['heading_residuals'] = heading_residuals_normalized * (np.pi/NUM_HEADING_BIN) # BxNUM_HEADING_BIN

    size_scores = tf.slice(output, [0,3+NUM_HEADING_BIN*2], [-1,NUM_SIZE_CLUSTER]) # BxNUM_SIZE_CLUSTER
    size_residuals_normalized = tf.slice(output, [0,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER], [-1,NUM_SIZE_CLUSTER*3])
    size_residuals_normalized = tf.reshape(size_residuals_normalized, [batch_size, NUM_SIZE_CLUSTER, 3]) # BxNUM_SIZE_CLUSTERx3
    end_points['size_scores'] = size_scores
    end_points['size_residuals_normalized'] = size_residuals_normalized
    end_points['size_residuals'] = size_residuals_normalized * tf.expand_dims(tf.constant(mean_size_arr, dtype=tf.float32), 0)

    return logits, end_points

def huber_loss(error, delta):
    abs_error = tf.abs(error)
    quadratic = tf.minimum(abs_error, delta)
    linear = (abs_error - quadratic)
    losses = 0.5 * quadratic**2 + delta * linear
    return tf.reduce_mean(losses)

# TODO: Test correctness of the loss...
def get_loss(logits, \
             mask_label, center_label, \
             heading_class_label, heading_residual_label, \
             size_class_label, size_residual_label, \
             end_points, reg_weight=0.001):
    """ logits: BxNxC,
        mask_label: BxN, """
    batch_size = logits.get_shape()[0].value
    mask_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=mask_label))
    tf.summary.scalar('3d mask loss', mask_loss)

    center_dist = tf.norm(center_label - end_points['center'], axis=-1)
    center_loss = huber_loss(center_dist, delta=2.0)
    tf.summary.scalar('center loss', center_loss)

    stage1_center_dist = tf.norm(center_label - end_points['stage1_center'], axis=-1)
    stage1_center_loss = huber_loss(stage1_center_dist, delta=1.0)
    tf.summary.scalar('stage1 center loss', stage1_center_loss)

    heading_class_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=end_points['heading_scores'], labels=heading_class_label))
    tf.summary.scalar('heading class loss', heading_class_loss)

    tmp = tf.one_hot(heading_class_label, depth=NUM_HEADING_BIN, on_value=1, off_value=0, axis=-1) # BxNUM_HEADING_BIN
    print tmp
    heading_residual_normalized_label = heading_residual_label / (np.pi/NUM_HEADING_BIN)
    heading_residual_normalized_loss = huber_loss(tf.reduce_sum(end_points['heading_residuals_normalized']*tf.to_float(tmp), axis=1) - heading_residual_normalized_label, delta=1.0)
    print heading_residual_normalized_loss
    tf.summary.scalar('heading residual normalized loss', heading_residual_normalized_loss)

    size_class_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=end_points['size_scores'], labels=size_class_label))
    tf.summary.scalar('size class loss', size_class_loss)

    tmp2 = tf.one_hot(size_class_label, depth=NUM_SIZE_CLUSTER, on_value=1, off_value=0, axis=-1) # BxNUM_SIZE_CLUSTER
    tmp2_tiled = tf.tile(tf.expand_dims(tf.to_float(tmp2), -1), [1,1,3]) # BxNUM_SIZE_CLUSTERx3
    predicted_size_residual_normalized = tf.reduce_sum(end_points['size_residuals_normalized']*tmp2_tiled, axis=[1]) # Bx3

    tmp3 = tf.expand_dims(tf.constant(mean_size_arr, dtype=tf.float32),0) # 1xNUM_SIZE_CLUSTERx3
    mean_size_label = tf.reduce_sum(tmp2_tiled * tmp3, axis=[1]) # Bx3
    size_residual_label_normalized = size_residual_label / mean_size_label

    size_normalized_dist = tf.norm(size_residual_label_normalized - predicted_size_residual_normalized, axis=-1)
    size_residual_normalized_loss = huber_loss(size_normalized_dist, delta=1.0)
    print size_residual_normalized_loss
    tf.summary.scalar('size residual normalized loss', size_residual_normalized_loss)

    # Compute IOU 3D
    iou2ds, iou3ds = tf.py_func(compute_box3d_iou, [end_points['center'], end_points['heading_scores'], end_points['heading_residuals'], end_points['size_scores'], end_points['size_residuals'], center_label, heading_class_label, heading_residual_label, size_class_label, size_residual_label], [tf.float32, tf.float32])
    tf.summary.scalar('iou_2d', tf.reduce_mean(iou2ds))
    tf.summary.scalar('iou_3d', tf.reduce_mean(iou3ds))

    end_points['iou2ds'] = iou2ds
    end_points['iou3ds'] = iou3ds

    # Compute BOX3D corners
    corners_3d = get_box3d_corners(end_points['center'], end_points['heading_residuals'], end_points['size_residuals']) # (B,NH,NS,8,3)
    gt_mask = tf.tile(tf.expand_dims(tmp, 2), [1,1,NUM_SIZE_CLUSTER]) * tf.tile(tf.expand_dims(tmp2,1), [1,NUM_HEADING_BIN,1]) # (B,NH,NS)
    corners_3d_pred = tf.reduce_sum(tf.to_float(tf.expand_dims(tf.expand_dims(gt_mask,-1),-1))*corners_3d, axis=[1,2]) # (B,8,3)

    heading_bin_centers = tf.constant(np.arange(0,2*np.pi,2*np.pi/NUM_HEADING_BIN), dtype=tf.float32) # (NH,)
    heading_label = tf.expand_dims(heading_residual_label,1) + tf.expand_dims(heading_bin_centers, 0) # (B,NH)
    heading_label = tf.reduce_sum(tf.to_float(tmp)*heading_label, 1)
    mean_sizes = tf.expand_dims(tf.constant(mean_size_arr, dtype=tf.float32), 0) # (1,NS,3)
    size_label = mean_sizes + tf.expand_dims(size_residual_label, 1) # (1,NS,3) + (B,1,3) = (B,NS,3)
    size_label = tf.reduce_sum(tf.expand_dims(tf.to_float(tmp2),-1)*size_label, axis=[1]) # (B,3)
    corners_3d_gt = get_box3d_corners_helper(center_label, heading_label, size_label) # (B,8,3)
    corners_3d_gt_flip = get_box3d_corners_helper(center_label, heading_label+np.pi, size_label) # (B,8,3)

    corners_dist = tf.minimum(tf.norm(corners_3d_pred - corners_3d_gt, axis=-1), tf.norm(corners_3d_pred - corners_3d_gt_flip, axis=-1))
    print "Corners dist: ", corners_dist
    corners_loss = huber_loss(corners_dist, delta=1.0)
    tf.summary.scalar('corners loss', corners_loss)

    return mask_loss + (center_loss + heading_class_loss + size_class_loss + heading_residual_normalized_loss*20 + size_residual_normalized_loss*20 + stage1_center_loss)*0.1 + corners_loss

if __name__=='__main__':
    with tf.Graph().as_default():
        inputs = tf.zeros((32,1024,4))
        outputs = get_model(inputs, tf.ones((32,3)), tf.constant(True))
        print outputs
        loss = get_loss(outputs[0], tf.zeros((32,1024),dtype=tf.int32), tf.zeros((32,3)), tf.zeros((32,),dtype=tf.int32), tf.zeros((32,)), tf.zeros((32,),dtype=tf.int32), tf.zeros((32,3)), outputs[1])
        print loss


"""
device.py - A module for defining attributes in the device..
"""

#from fqc.util import get_unitary
import numpy as np
import scipy.linalg as la
from scipy.special import factorial
import os,sys,inspect
# import h5py
import matplotlib.pyplot as plt
import random as rd
import time
import re, math
from datetime import datetime
from ..util import get_grid_coupling_list, get_connectivity_graph, get_crosstalk_graph

GATETIMES = {'unitary': 55,'rz': 30, 'z':30, 'u1': 30, 's': 30, 't': 30, 'rx': 30, 'x': 30, 'u2': 30, 'ry': 30, 'y': 30, 'u3': 30, 'h': 30, 'measure': 0.0, 'barrier': 0.0} # all in ns

class Device(object):
    """
    Fields:
    """
    def __init__(self, topology, qubits, omega_max, delta_int, delta_ext, delta_park, cqq=0.019, alpha=-0.2, EJS=8, EJL=20, EC=0.5, flux_sigma=0.0, error_1q_gate=0.001, d=1):
        if (topology == None):
            self.topology = 'grid'
        else:
            self.topology = topology
        self.qubits = qubits
        self.side_length = int(np.sqrt(qubits)) # only for square grid
        self.omega_max = omega_max
        self.delta_int = delta_int
        self.delta_ext = delta_ext
        self.delta_park = delta_park
        self.omega_min = omega_max - delta_int - delta_ext - delta_park
        self.cqq = cqq
        self.alpha = alpha
        self.ejs = EJS
        self.ejl = EJL
        self.ec = EC
        self.flux_sigma = flux_sigma
        self.g_connect = None
        self.coupling = None
        self.g_xtalk = None
        self.d_xtalk = d
        self.gate_times = GATETIMES
        self.error_1q_gate = error_1q_gate

    def build_graph(self, param=None):
        if self.topology=='grid':
            self.g_connect = get_connectivity_graph(self.qubits, self.topology)
            self.coupling = get_grid_coupling_list(self.side_length, self.side_length)
            self.g_xtalk = get_crosstalk_graph(self.g_connect, self.topology, self.d_xtalk)
        else:
            self.g_connect = get_connectivity_graph(self.qubits, self.topology, param)
            self.coupling = self.g_connect.edges()
            self.g_xtalk = get_crosstalk_graph(self.g_connect, self.topology, self.d_xtalk)

class Sycamore_device(object):
    def __init__(self, device, size, res_coupling=0.0):
        if size != device.qubits:
            print("Warning: device size inconsistent. Device: " + str(device.qubits) + ", Sycamore_device: " + str(size))
        if size not in [4,9,16,25]:
            raise Exception("Wrong device size")
        if size == 4:
            self.int_freqs = {(0,1):6.619, (0,2):6.65, (1,3):6.657, (2,3):6.677}
            self.park_freqs = {0:6.605,1:6.638,2:6.694,3:6.681}
        elif size == 9:
            self.park_freqs = {0:6.605,1:6.638,2:6.565,3:6.694,4:6.681,5:6.601,6:6.643,7:6.621,8:6.646}
            self.int_freqs = {(0,1):6.619, (1,2):6.601, (3,4):6.677, (4,5):6.635, (6,7):6.631, (7,8):6.646, (0,3):6.65, (1,4):6.657, (2,5):6.585, (3,6):6.667, (4,7):6.645,(5,8):6.646}
        elif size == 16:
            self.park_freqs = {0:6.605,1:6.638,2:6.565,3:6.555,4:6.694,5:6.681,6:6.601,7:6.626,8:6.643,9:6.621,10:6.646,11:6.657,12:6.712,13:6.671,14:6.586,15:6.623}
            self.int_freqs = {(0,1):6.619,(1,2):6.601,(2,3):6.565,(4,5):6.677,(5,6):6.635,(6,7):6.595}
            self.int_freqs[(8,9)] = 6.631
            self.int_freqs[(9,10)] = 6.646
            self.int_freqs[(10,11)] = 6.646
            self.int_freqs[(12,13)] = 6.69
            self.int_freqs[(13,14)] = 6.631
            self.int_freqs[(14,15)] = 6.623
            self.int_freqs[(0,4)] = 6.65
            self.int_freqs[(1,5)] = 6.657
            self.int_freqs[(2,6)] = 6.585
            self.int_freqs[(3,7)] = 6.592
            self.int_freqs[(4,8)] = 6.667
            self.int_freqs[(5,9)] = 6.645
            self.int_freqs[(6,10)] = 6.646
            self.int_freqs[(7,11)] = 6.642
            self.int_freqs[(8,12)] = 6.68
            self.int_freqs[(9,13)] = 6.645
            self.int_freqs[(10,14)] = 6.646
            self.int_freqs[(11,15)] = 6.633
        else: # size == 25
            self.park_freqs = {0:6.612,1:6.571,2:6.605,3:6.638,4:6.565,5:6.687,6:6.661,7:6.694,8:6.681,9:6.601,10:6.634,11:6.628,12:6.643,13:6.621,14:6.646,15:6.707,16:6.665,17:6.712,18:6.671,19:6.586,20:6.775,21:6.734,22:6.766,23:6.729,24:6.594}
            self.int_freqs = {(0,1):6.592,(1,2):6.589,(2,3):6.619,(3,4):6.601}
            self.int_freqs[(5,6)] = 6.675
            self.int_freqs[(6,7)] = 6.678
            self.int_freqs[(7,8)] = 6.677
            self.int_freqs[(8,9)] = 6.635
            self.int_freqs[(10,11)] = 6.628
            self.int_freqs[(11,12)] = 6.632
            self.int_freqs[(12,13)] = 6.631
            self.int_freqs[(13,14)] = 6.646
            self.int_freqs[(15,16)] = 6.693
            self.int_freqs[(16,17)] = 6.690
            self.int_freqs[(17,18)] = 6.690
            self.int_freqs[(18,19)] = 6.631
            self.int_freqs[(20,21)] = 6.753
            self.int_freqs[(21,22)] = 6.766
            self.int_freqs[(22,23)] = 6.743
            self.int_freqs[(23,24)] = 6.594

            self.int_freqs[(0,5)] = 6.648
            self.int_freqs[(1,6)] = 6.617
            self.int_freqs[(2,7)] = 6.650
            self.int_freqs[(3,8)] = 6.657
            self.int_freqs[(4,9)] = 6.585
            self.int_freqs[(5,10)] = 6.660
            self.int_freqs[(6,11)] = 6.640
            self.int_freqs[(7,12)] = 6.667
            self.int_freqs[(8,13)] = 6.645
            self.int_freqs[(9,14)] = 6.646
            self.int_freqs[(10,15)] = 6.670
            self.int_freqs[(11,16)] = 6.646
            self.int_freqs[(12,17)] = 6.680
            self.int_freqs[(13,18)] = 6.645
            self.int_freqs[(14,19)] = 6.646
            self.int_freqs[(15,20)] = 6.741
            self.int_freqs[(16,21)] = 6.704
            self.int_freqs[(17,22)] = 6.712
            self.int_freqs[(18,23)] = 6.703
            self.int_freqs[(19,24)] = 6.594

        self.res_coupling = res_coupling
        omega_max = max(self.int_freqs.values())
        omega_min = min(self.int_freqs.values())
        if omega_max > device.omega_max:
            print("Warning: max freq inconsistent. Device: " + str(device.omega_max) + ", Sycamore_device: " + str(omega_max))
            self.omega_max = omega_max
        if omega_min < device.omega_min:
            print("Warning: min freq inconsistent. Device: " + str(device.omega_min) + ", Sycamore_device: " + str(omega_min))
            self.omega_min = omega_min

    def get_park_freq(self,q):
        if q not in self.park_freqs:
            raise Exception("Wrong qubit number")
        return self.park_freqs[q]

    def get_itr_freq(self,q1,q2):
        if (q1,q2) in self.int_freqs:
            return self.int_freqs[(q1,q2)]
        elif (q2,q1) in self.int_freqs:
            return self.int_freqs[(q1,q2)]
        else:
            raise Exception("Coupling not found")


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Interface to the Fermi and HESS catalogs.
"""
from __future__ import absolute_import, division, print_function, unicode_literals
import logging
import numpy as np
from astropy.table import Table
from astropy.units import Unit, Quantity
from ..spectrum import compute_differential_flux_points

__all__ = [
    'SEDComponent',
    'SED',
    'cube_sed',
    'add_spec',
]

log = logging.getLogger(__name__)

MeV_to_GeV = Unit('MeV').to(Unit('GeV'))
MeV_to_erg = Unit('MeV').to(Unit('erg'))

class SEDComponent(object):
    """Uniform interface to SED components for the SED class
    """

    def __init__(self, name='', model=None, points=None):
        """
        @param name: str
        @type model: spec.spectrum.Spectrum
        @type points: spec.data.FluxPoints"""
        self.name = name
        self.model = model
        self.points = points

    def plot(self, model=True, points=True, butterfly=True):
        if butterfly:
            self.plot_butterfly()
        if model:
            self.plot_model()
        if points:
            self.plot_points()

    def plot_model(self):
        import matplotlib.pyplot as plt
        if self.model is None:
            log.warning('{0}: No model available.'.format(self.name))
            return
        x, y = self.model.points(power=2)
        plt.plot(x * MeV_to_GeV, y * MeV_to_erg, label=self.name)

    def plot_points(self, color='black', markerfacecolor='black'):
        import matplotlib.pyplot as plt
        if self.points is None:
            log.warning('{0}: No points available.'.format(self.name))
            return
        # Note: We plot each point individually because anyway
        # upper limits have to be plotted differently which I
        # think is not possible because the marker argument doesn't
        # take arrays.
        for ii in range(len(self.points)):
            x, exl, exh, y, eyl, eyh, ul = self.points[ii]
            if ul:
                marker = ' '
                lolims = True
                eyl, eyh = 5 / 10. * y, 0
            else:
                marker = 'o'
                lolims = False
            plt.plot(MeV_to_GeV * x, y)
            plt.errorbar(x, y, [eyl, eyh], [exl, exh], lolims=lolims,
                         marker=marker, color=color,
                         markerfacecolor=markerfacecolor,
                         label=self.name)

    def plot_butterfly(self):
        pass

class SED(list):
    """Class to plot GeV -- TeV SEDs

    Internally the same units as in the Fermi catalog are used:
    - Energies in MeV
    - Flux densities in cm^-2 s^-2 MeV^-1
    - Fluxes in cm^-2 s^-1
    - Energy fluxes in erg cm^-2 s^-1"""
    """
    def add_Fermi(self, name):
        try:
            self._fermi
            self.append(self._fermi.sed_component(name))
        except
        component = catalog.sed_component(name)
        self.append(component)
    """

    def add(self, names, catalogs):
        for name in names:
            for catalog in catalogs:
                try:
                    component = catalog.sed_component(name)
                    self.append(component)
                    log.info('%s found in %s',
                             name, catalog.table.table_name)
                except ValueError as e:
                    log.warning(e)
                    log.warning('%s not found in %s',
                                name, catalog.table.table_name)
                    pass

    def plot(self, filename='sed.png', xlim=(8e-2, 2e5), ylim=(1e-14, 1e-8)):
        import matplotlib.pyplot as plt
        plt.figure()
        plt.ylabel(r'E$^2$ dF/DE (erg cm$^{-2}$ s$^{-1}$)')
        plt.xlabel('Energy (GeV)')
        plt.loglog()
        log.info('Plotting {0} components in SED'.format(len(self)))
        for component in self:
            component.plot()
        plt.xlim(xlim)
        plt.ylim(ylim)
        plt.legend()
        log.info('Writing {0}'.format(filename))
        plt.savefig(filename)

    def add_component(self, catalog_format, catalog_name,
                      object_name, plot_pivot=False, **ecpl_params):
        """ Read necessary parameters from FITS file and plot butterfly

        Parameters:
        catalog_format = 'hess', 'fermi'
        catalog_name = FITS file name
        object_name  = object name string in 'name' column

        Note: Since every catalog has columns with different
        names and units, a general SED plotting is not possible.
        Instead for each catalog type a handler function that
        deals converts to a standard format is called.

        @todo: Possibly pass plotting parameters along here by
        appending them to the ecpl_params dictionary
        -> I don't think this works at the moment!!!"""
        from atpy import Table
        # Get the catalog from file and initialize some things
        self.catalog_format = catalog_format
        self.catalog_name = catalog_name
        self.object_name = object_name
        self.catalog = Table(catalog_name).data
        # Build a dictionary of parameters needed for the plot
        self.ecpl_params = ecpl_params
        self.get_ecpl_params()
        # Plot curve
        self.plot_ecpl(plot_pivot=plot_pivot, **ecpl_params)
        # Plot points if present
        if self.plot_points is not None:
            # Get the values needed for plotting
            e = self.plot_points[0]
            f = self.plot_points[1]
            f_err = self.plot_points[2]
            e_err = self.plot_points[3]
            is_ul = self.plot_points[4]
            for ii in range(e.size):
                self.plot_point(e[ii], f[ii],
                                f_err=f_err[ii],
                                e_err=[[e_err[0][ii]], [e_err[1][ii]]],
                                ul=is_ul[ii])
            # Remove so that it doesn't get plotted again.
            self.plot_points = None

    def get_ecpl_params(self):
        """Build self.ecpl_params dictionary
        by parsing one of the supported catalogs"""
        if self.catalog_format == 'hess':
            self.get_ecpl_params_hess_cat()
        elif self.catalog_format == 'fermi':
            self.get_ecpl_params_fermi_cat()
        # Change numpy types to regular types
        # and replace nan values with 0
        for key, value in self.ecpl_params.items():
            if isinstance(value, np.float32):
                value = float(value)
            if isinstance(value, np.int16):
                value = int(value)

    def get_ecpl_params_fermi_cat(self):
        """ Build self.ecpl_params dictionary from Fermi catalog fields """
        i = self.find_object_index('source_name')
        # Set all plot parameters:
        self.ecpl_params['e_pivot'] = self.catalog.field('Pivot_Energy')[i]
        self.ecpl_params['e_min'] = 1e2
        self.ecpl_params['e_max'] = 1e5
        self.ecpl_params['e_cut'] = 0.0
        self.ecpl_params['e_cut_err'] = 0.0
        self.ecpl_params['e_scale'] = 1
        self.ecpl_params['norm'] = self.catalog.field('Flux_Density')[i]
        self.ecpl_params['norm_err'] = self.catalog.field('Unc_Flux_Density')[i]
        self.ecpl_params['norm_scale'] = 1
        self.ecpl_params['index'] = self.catalog.field('Spectral_Index')[i]
        self.ecpl_params['index_err'] = self.catalog.field('Unc_Spectral_Index')[i]
        self.ecpl_params['color'] = 'green'
        self.ecpl_params['butterfly'] = True
        # Set flux point data
        self.plot_points = self.get_flux_points_fermi(i)
        # Add text label
        fmt = '%s\n%s, %s\n' + \
              r'S = %3.1f, C = %3.1f, $\Gamma = %1.2f \pm %1.2f$'
        values = (self.object_name,
                  self.catalog.field('class1')[i],
                  self.catalog.field('assoc1')[i],
                  self.catalog.field('signif_avg')[i],
                  self.catalog.field('curvature_index')[i],
                  self.catalog.field('spectral_index')[i],
                  self.catalog.field('unc_spectral_index')[i]
                  )
        self.ax.text(0.05, 0.95, fmt % values,
                     horizontalalignment='left',
                     verticalalignment='top',
                     transform=self.ax.transAxes)

def add_spec(frame, model, xlim, npoints=100, **plot_params):
    """Add a spectral component to a frame.

    frame = matplotlib.Axes object
    model = [function, parameters, constants]
    xlim  = [xmin, xmax]"""
    # Unpack model
    f, p, c = model
    # Compute x and y values
    logx = np.linspace(np.log10(xlim[0]), np.log10(xlim[1]), npoints)
    x = 10 ** logx
    y = f(p, c, x)
    frame.plot(x, y, **plot_params)

def add_crab(ax):
    """Add the Fermi and HESS Crab SED to test scaling."""
    pass
    # The HESS butterfly
    # Note: The HESS catalog contains energies in TeV and flux norm in 1e-12 cm^-2 s^-1 TeV^-1
    """
    add_sed_component(ax, e0 = 1, e1 = 1e-2, e2 = 1e2,
                      norm = 10, norm_err = 0, index = 2., index_err = 0.0,
                      e_scale = 1e12, norm_scale = 1e-12 * 1e-12, e_cut = 10, e_cut_err = 3,
                      color='b', butterfly = True)
    # The Fermi butterfly
    add__sed_component(ax, e0 = 494, e1 = 1e2, e2 = 1e6,
                       norm = 1e-9, norm_err = 6.7e-11, index = 2.3, index_err = 0.1,
                       e_scale = 1e6, norm_scale = 1e-6, color='g', butterfly = True)
    # Add published fermi result
    """

def cube_sed(cube, mask=None, flux_type='differential', counts=None,
             errors=False, standard_error=0.1, spectral_index=2.3):
    """Creates SED from SkyCube within given lat and lon range.

    Parameters
    ----------
    cube : `~gammapy.data.SkyCube`
        Spectral cube of either differential or integral fluxes (specified
        with flux_type)
    mask : array_like, optional
        2D mask array, matching spatial dimensions of input cube.
        A mask value of True indicates a value that should be ignored,
        while a mask value of False indicates a valid value.
    flux_type : {'differential', 'integral'}
        Specify whether input cube includes differential or integral fluxes.
    counts :  `~gammapy.data.SkyCube`, optional
        Counts cube to allow Poisson errors to be calculated. If not provided,
        a standard_error should be provided, or zero errors will be returned.
    errors : bool
        If True, computes errors, if possible, according to provided inputs.
        If False (default), returns all errors as zero.
    standard_error : float
        If counts cube not provided, but error values required, this specifies
        a standard fractional error to be applied to values. Default = 0.1.
    spectral_index : float
        If integral flux is provided, this is used to calculate differential
        fluxes and energies (according to the Lafferty & Wyatt model-based
        method, assuming a power-law model).

    Returns
    -------
    table : `~astropy.table.Table`
        A spectral energy table of energies, differential fluxes and
        differential flux errors. Units as those input.
    """

    values = []
    for i in np.arange(cube.data.shape[0]):
        if mask is None:
            bin = cube.data[i].sum()
        else:
            bin = cube.data[i][mask].sum()
        values.append(bin.value)
    values = np.array(values)

    if errors:
        if counts is None:
            # Counts cube required to calculate poisson errors
            errors = np.ones_like(values) * standard_error
        else:
            errors = []
            for i in np.arange(counts.data.shape[0]):
                if mask is None:
                    bin = counts.data[i].sum()
                else:
                    bin = counts.data[i][mask].sum()
                r_error = 1. / (np.sqrt(bin.value))
                errors.append(r_error)
            errors = np.array(errors)
    else:
        errors = np.zeros_like(values)

    if flux_type == 'differential':
        energy = cube.energy
        table = Table()
        table['ENERGY'] = energy
        table['DIFF_FLUX'] = Quantity(values, cube.data.unit)
        table['DIFF_FLUX_ERR_HI'] = Quantity(errors * values, cube.data.unit)
        table['DIFF_FLUX_ERR_LO'] = Quantity(-errors * values, cube.data.unit)

    elif flux_type == 'integral':

        emins = cube.energy[:-1]
        emaxs = cube.energy[1:]
        table = compute_differential_flux_points(x_method='lafferty',
                                                 y_method='power_law',
                                                 spectral_index=spectral_index,
                                                 energy_min=emins, energy_max=emaxs,
                                                 int_flux=values,
                                                 int_flux_err_hi=errors * values,
                                                 int_flux_err_lo=-errors * values)

    else:
        raise ValueError('Unknown flux_type: {0}'.format(flux_type))

    return table


import numpy as np
from KiMonETSim.conversion_functions import from_ev_to_au, from_ns_to_au

"""
COMENTARI INICIAL: Els estats (fonamental i excitats) s'indicaran per mitjà d'etiquetes tipus string. Els estats possibles
i les seves respectives referències hauran de venir indicades a l'inici del programa, e.g, quan s'inicialitzi la
molècula genèrica. Aquestes referències s'usaran per caracteritzar la variable estat i per simplicitat i eficiència
s'usaran com a claus dels diccionaris que contenguin informació sobre els excitons. A saber: state_energies,
relaxation_energies, exciton_energies, (REFERENTS A RATES...).

INICIALITZACIÓ DE LA MOLÈCULA GENÈRICA:
    -state_energies: diccionari amb les energies d'excitació de cada estat de la molècula:
                :key: energia de l'estat excitat
                :argument:  energia de l'estat
    -relaxation_energies: diccionari amb les energies de relaxació de cada estat. Per ara en prenem una fitxa i què depèn
    sols de cada estat.
                :key: energia de l'estat excitat
                :argument:  energia de relaxació de l'estat
    -transition_moment: moment de transició dipolar de la molècula. Donat com un vector (llista de 3 elements) en el
    sistema de referència de la molècula.

Per defecte venen donats:
    -characteristic_length: defineix les dimensions finites de la molècula. Aquesta s'aproxima com una línia, quadrat o
    cub i aquest paràmetre en defineix el costat. Alternativament podriem fer una aproximació esfèrica i que en sigui
    el radi.
    -coordinates. D'entrada la suposam en l'origen.
    -orientation. D'entrada la suposam orientada segons [1, 0, 0]. Aquest vector ve donat en un SR extern que
    anomenarem global.
    Aquests dos darrers paràmetres no són estrictament necessaris per estudiar la naturalesa del tipus de molècula.
    La classe inclou 2 mètodes per cada un d'aquests 2 darrers paràmetres. Un per inicialitzar-los (defineix de manera
    la posició/orientació com un 3-array) i un per cridar-los alhora d'operar (per assegurar que no s'alteren en el procés).
    Noms:
    initialize_coordinates(coordinates)                  initialize_orientation(orientation)
    molecular_coordinates()                              molecular_orientation()

Mètodes de la molècula:
Apart dels 4 ja comentats la classe molècula inclou:
    - electronic_state. Mètode que retorna l'estat electrònic de la molecula
    - get_relaxation_state_energy. Mètode que dóna l'energia de relaxació de l'estat en què es troba la molècula
    - change_state(new_state): Canvia l'estat de la molècula pel nou donat.
    - decay_rates: mètode que retorna un diccionari amb els possibles rates de decaïment {'decay process': rate}
    - get_transition_moment(reference_orientation). Necessita com argument un vector de referència. L'orientació de la
        molècula en la qual el moment de transició en el SR de la molècula i en el SR global coindideixen.
        Aleshores, donada aquesta referència i l'orientació de la molècula, aquest mètode fa un canvi de base i retorna
        el moment de transició dipolar en el SR global.
"""

class Molecule:

    def __init__(self,
                 state_energies,
                 reorganization_energies,
                 transition_moment,
                 dipole_moment_direction=[1, 0, 0],                 # unity vector
                 state='gs',
                 characteristic_length=10**(-8),
                 coordinates=[0, 0, 0],
                 orientation=[1, 0, 0]):                            # unity vector
        """
        :param states_energies: dictionary {'state': energy}
        :param state: sting of the name of the state
        The name of the state should coincide with some key of the dictionary in order to identify the state with
        its energy.
        :param reorganization_energies: dictionary {'state': relaxation energy of the state}
        Names of 'state' would be: g_s (ground state), s_1 (first singlet), t_1 (first triplet), etc.
        Energies should be given with eV.

        :param transition_moment: Dipole transition moment vector (3d). The vector is given in respect to the RS
        of the molecule. So for all molecules of a same type if will be equal.
        This dipole moment is given in atomic units.

        :param dipole_moment_direction: Gives a referene direction This reference direction is the molecular orientation
        in which the transition dipole moment in the molecular reference system coincides with the itself in the
        global reference system. Given by default as [1,0,0]

        :param characteristic_length: We consider a finite size molecule. The simplified shape of the molecule
        is longitudinal, squared or cubic and is defined with this characteristic length. Units: nm

        :param coordinates: 3d vector. Gives the position of the molecule in the system (in general the 0 position
        will coincide with the center of the distribution). Units: nm. If the system has less than 3 dimensions,
        the extra coordinates will be taken as 0.
        :param orientation: 3d unit vector. Gives the orientation of the molecule in the global reference system.
        """

        self.state_energies = state_energies
        self.state = state
        self.reorganization_energies = reorganization_energies
        self.transition_moment = np.array(transition_moment)
        self.dipole_moment_direction = np.array(dipole_moment_direction)        # unity vector
        self.characteristic_length = characteristic_length
        self.coordinates = np.array(coordinates)
        self.orientation = np.array(orientation)                                # unity vector

    def initialize_coordinates(self, coordinate_list):
        """
        :param coordinate_list: List [x, y, z] with the coordinates of the molecule. Units: nm
        Changes self.coordinates to this new position. Format: numpy array.
        """
        self.coordinates = np.array(coordinate_list)

    def molecular_coordinates(self):
        """
        :return: Array with the molecular coordinates.
        """
        return self.coordinates

    def initialize_orientation(self, orientation):
        """
        :param orientation: list with the coordinates of the orientation vector
        Changes self.orientation to this new orientation. Format: numpy array
        """
        self.orientation = np.array(orientation)

    def molecular_orientation(self):
        """
        :return: Array with the molecular orientation
        """
        return self.orientation

    def get_reorganization_state_energy(self):
        return self.reorganization_energies[self.state]

    def electronic_state(self):
        """
        :return: the electronic state of the molecule
        """
        return self.state

    def set_state(self, new_state):
        """
        :param new_state:
        No return method. Only changes the molecular state when the exciton is transferred.
        """
        self.state = new_state

    def desexcitation_energies(self):
        """
        IS NOT USED (19/08/2019).
        Given an electronic state, calculates the possible desexcitation energy. Generates and sorts
        a list with the energies, then calculates the possible desexcitation energies (the energy difference
        between the given state and the less energetic states).
        :return: Dictionary with the decay processes as key, e.g. 'State1_to_State0', and the energy as argument
        """
        desexcitations = {}

        for state_key in self.state_energies:
            if self.state_energies[self.state] > self.state_energies[state_key]:
                decay_process = 'from_'+self.state+'_to_'+state_key
                energy_gap = self.state_energies[self.state] - self.state_energies[state_key]
                desexcitations[decay_process] = energy_gap

        return desexcitations

    def decay_rates(self):
        """
        :return: A list of two elements: list of the possible decay processes and another with the respective rates
        for a given electronic state.

        More if(s) entrances shall be added if more electronic states are considered.
        """

        decay_rates = {}

        if self.state == 'gs':
            process = 'No decay for a molecule at the ground state.'
            print(process)
            decay_rates[process] = 0

        elif self.state == 's1':

            desexcitation_energy = self.state_energies[self.state] - self.state_energies['gs']      # energy in eV
            desexcitation_energy = from_ev_to_au(desexcitation_energy, 'direct')                    # energy in a.u.

            u = np.linalg.norm(self.transition_moment)              # transition moment norm.
            c = 137                                                 # light speed in atomic units

            rate = 4 * desexcitation_energy**3 * u**2 / (3 * c**3)
            decay_process = 'Singlet_radiative_decay'
            # for a first singlet state only radiative decay is considered.

            decay_rates[decay_process] = from_ns_to_au(rate, 'direct')

        return decay_rates

    def get_transition_moment(self):
        """
        This method computes a basis transformation in order to get the transition dipole moment in a global
        reference coordinate system. This system is described by the dipole_moment_direction
        When the molecule is orientated in this way the transition_moment in the RS system of the molecule
        and in the global RS coincides. Else, we can say that the molecule is rotated and the rotation angle
        is given by the inner product between the reference_orientation and molecule_orientation vectors
        :return: The t.d.m in the global reference system
        """

        # compute the director cosinus of the rotation (inner product)
        cos_director = np.dot(self.dipole_moment_direction, self.orientation)

        # construct the rotation matrix
        sin_director = np.sqrt(1-cos_director**2)

        rotation_matrix = np.array([[cos_director, -sin_director, 0],
                                    [sin_director, cos_director, 0],
                                    [0,                 0,       0]])

        # basis transform (matrix product)
        return np.dot(rotation_matrix, self.transition_moment)


license="""
   Copyright (C) 2015 James Annis

   This program is free software; you can redistribute it and/or modify it
   under the terms of version 3 of the GNU General Public License as
   published by the Free Software Foundation.

   More to the points- this code is science code: buggy, barely working,
   with little or no documentation. Science code in the the alpine fast
   & light style. (Note the rate at which people who stand on the
   summit of K2 to successfully make it down.)

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
"""

import numpy as np
import decam2hp

#
# Example usage
#
#    obs.resetTime(mjd+time)
#    obs.limitMag("i", exposure=90)
#
#    sm=sourceProb.map(obs);
#    models_at_t = modelsAtTimeT (models, time)
#    abs_mag = models_at_t[0]
#    sm.modelAbsoluteMagnitude = abs_mag
#    sm.searchDistance = np.array([distance,])
#    sm.calculateProb()
#
#    raHex, decHex, idHex, hexVals, rank = cutAndHexalate (obs, sm, raHexen, decHexen)

#
# Given an obs object and a sm object, along with a map of hex centers,
# this routine returns the sum of the probability inside of DECam camera outlines
# for all hexes in the map, along with a rank of the hexes sorted from large to small.
#
# domain knowlege- the blanco cuts keep everything inside an hour angle range
def hexalateNHexes (obs, sm, nHexes, allskyDesHexes) :
    raHexen, decHexen, idHexen, hexVals, rank = cutAndHexalate(obs, sm, allskyDesHexes)
    raHexen = raHexen[0:nHexes]
    decHexen = decHexen[0:nHexes]
    idHexen = idHexen[0:nHexes]
    hexVals = hexVals[0:nHexes]
    rank = rank[0:nHexes]
    return raHexen, decHexen, idHexen, hexVals, rank

# be aware that while most of my ra,dec are in degrees,
# those in obs and sm are in radians
def cutAndHexalate (obs, sm, camera, hexFile) :
    #allskyDesHexes="../data/all-sky-hexCenters-"+camera+".txt"
    raHexen, decHexen, idHexen = getHexCenters(hexFile)
    raHexen, decHexen, idHexen, hexVals, rank = cutAndHexalateOnRaDec(
        obs, sm, raHexen, decHexen, idHexen, camera)
    return raHexen, decHexen, idHexen, hexVals, rank

def cutAndHexalateOnRaDec (obs, sm, raHexen, decHexen, idHexen, tree, camera, cutProbs=False) :
    #cutProbs calls without Overlap
    verbose = False
    obsHourAngle = obs.ha*360./(2*np.pi)
    obsRa        = obs.ra*360./(2*np.pi)
    obsDec       = obs.dec*360./(2*np.pi)
    # based on blanco horizen limits  (may not be needed with tree)
    ix = (abs(obsHourAngle) <= 83. ) & (obsDec < 43.)
    ix2 = decHexen < 43.

    probabilities = obs.map*sm.probMap
    if verbose  :
        print "\t cutAndHexalate probabilities sum",probabilities.sum()

    hexVals = np.zeros(raHexen.size)
    if not cutProbs:
        hexVals[ix2] = decam2hp.hexalateMap(obsRa,obsDec, probabilities, tree,
                                            raHexen[ix2], decHexen[ix2], camera, verbose=False)
    else:
        hexVals[ix2] = decam2hp.hexalateMapWithoutOverlap(obsRa,obsDec, probabilities, tree,
                                            raHexen[ix2], decHexen[ix2], camera, verbose=False)
    if verbose  :
        print "hexVals max", hexVals.max()
    rank=np.argsort(hexVals);
    rank = rank[::-1];# sort from large to small by flipping natural argsort order

    return raHexen, decHexen, idHexen, hexVals, rank

def getHexCenters (hexFile) :
    #allskyDesHexes="../data/all-sky-hexCenters-"+camera+".txt"
    ra,dec = np.genfromtxt(hexFile, unpack=True, \
        usecols=(0,1),comments="#")
    hex_id = getHexId(ra,dec)
    return ra,dec,hex_id

#
# a modified DES convention: just rounded ra,dec with + or - sign
# of type "24-31", "179+10"
def getHexId(ra, dec) :
    intra = np.round(ra)
    intdec = np.abs(np.round(dec));
    sign=np.full(ra.size, "+",dtype="str");
    ix = dec < 0; sign[ix]="-";
    id = np.array([],dtype="str")
    for i in range(0,ra.size) :
        name = str(np.int(intra[i])) + sign[i] + str(np.int(intdec[i]))
        id = np.append(id, name)
    return id


#!/usr/bin/env python
"""
feffdat  provides the following function related to
reading and dealing with Feff.data files in larch:

  path1 = read_feffdat('feffNNNN.dat')

returns a Feff Group -- a special variation of a Group -- for
the path represented by the feffNNNN.dat

  group  = ff2chi(paths)

creates a group that contains the chi(k) for the sum of paths.
"""
import numpy as np
from copy import deepcopy
from scipy.interpolate import UnivariateSpline
from lmfit import Parameters, Parameter
from lmfit.printfuncs import gformat

from xraydb import atomic_mass, atomic_symbol

from larch import Group, isNamedClass
from larch.utils.strutils import fix_varname, b32hash
from larch.fitting import group2params, isParameter, param_value

from .xafsutils import ETOK, set_xafsGroup
from .sigma2_models import add_sigma2funcs

SMALL_ENERGY = 1.e-6

class FeffDatFile(Group):
    def __init__(self, filename,  **kws):
        kwargs = dict(name='feff.dat: %s' % filename)
        kwargs.update(kws)
        Group.__init__(self,  **kwargs)
        self._read(filename)

    def __repr__(self):
        if self.filename is not None:
            return '<Feff.dat File Group: %s>' % self.filename
        return '<Feff.dat File Group (empty)>'

    def __copy__(self):
        return FeffDatFile(filename=self.filename)

    def __deepcopy__(self, memo):
        return FeffDatFile(filename=self.filename)

    @property
    def reff(self): return self.__reff__

    @reff.setter
    def reff(self, val):     pass

    @property
    def nleg(self): return self.__nleg__

    @nleg.setter
    def nleg(self, val):     pass

    @property
    def rmass(self):
        """reduced mass for a path"""
        if self.__rmass is None:
            rmass = 0
            for atsym, iz, ipot, amass, x, y, z in self.geom:
                rmass += 1.0/max(1., amass)
            self.__rmass = 1./rmass
        return self.__rmass

    @rmass.setter
    def rmass(self, val):     pass

    def _read(self, filename):
        try:
            with open(filename, 'r') as fh:
                lines = fh.readlines()
        except:
            print( 'Error reading file %s ' % filename)
            return
        self.filename = filename
        mode = 'header'
        self.potentials, self.geom = [], []
        data = []
        pcounter = 0
        iline = 0
        for line in lines:
            iline += 1
            line = line[:-1]
            if line.startswith('#'): line = line[1:]
            line = line.strip()
            if iline == 1:
                self.title = line[:64].strip()
                self.version = line[64:].strip()
                continue
            if line.startswith('k') and line.endswith('real[p]@#'):
                mode = 'arrays'
                continue
            elif '----' in line[2:10]:
                mode = 'path'
                continue
            #
            if (mode == 'header' and
                line.startswith('Abs') or line.startswith('Pot')):
                words = line.replace('=', ' ').split()
                ipot, z, rmt, rnm = (0, 0, 0, 0)
                words.pop(0)
                if line.startswith('Pot'):
                    ipot = int(words.pop(0))
                iz = int(words[1])
                rmt = float(words[3])
                rnm = float(words[5])
                self.potentials.append((ipot, iz, rmt, rnm))
            elif mode == 'header' and line.startswith('Gam_ch'):
                words  = line.replace('=', ' ').split(' ', 2)
                self.gam_ch = float(words[1])
                self.exch   = words[2]
            elif mode == 'header' and line.startswith('Mu'):
                words  = line.replace('=', ' ').split()
                self.mu = float(words[1])
                self.kf = float(words[3])
                self.vint = float(words[5])
                self.rs_int= float(words[7])
            elif mode == 'path':
                pcounter += 1
                if pcounter == 1:
                    w = [float(x) for x in line.split()[:5]]
                    self.__nleg__ = int(w.pop(0))
                    self.degen, self.__reff__, self.rnorman, self.edge = w
                elif pcounter > 2:
                    words = line.split()
                    xyz = [float(x) for x in words[:3]]
                    ipot = int(words[3])
                    iz   = int(words[4])
                    if len(words) > 5:
                        lab = words[5]
                    else:
                        lab = atomic_symbol(iz)
                    amass = atomic_mass(iz)
                    geom = [lab, iz, ipot, amass] + xyz
                    self.geom.append(tuple(geom))
            elif mode == 'arrays':
                d = np.array([float(x) for x in line.split()])
                if len(d) == 7:
                    data.append(d)
        data = np.array(data).transpose()
        self.k        = data[0]
        self.real_phc = data[1]
        self.mag_feff = data[2]
        self.pha_feff = data[3]
        self.red_fact = data[4]
        self.lam = data[5]
        self.rep = data[6]
        self.pha = data[1] + data[3]
        self.amp = data[2] * data[4]
        self.__rmass = None  # reduced mass of path

PATH_PARS = ('degen', 's02', 'e0', 'ei', 'deltar', 'sigma2', 'third', 'fourth')

class FeffPathGroup(Group):
    def __init__(self, filename, label=None, s02=None, degen=None,
                 e0=None, ei=None, deltar=None, sigma2=None, third=None,
                 fourth=None, _larch=None, **kws):

        kwargs = dict(name='FeffPath: %s' % filename)
        kwargs.update(kws)
        Group.__init__(self, **kwargs)
        self.filename = filename
        self.params = None
        self.label = label
        self.spline_coefs = None

        self._feffdat = FeffDatFile(filename=filename)
        self.geom  = self._feffdat.geom
        def_degen  = self._feffdat.degen

        self.hashkey = self.__geom2label()
        self.label = label if label is not None else self.hashkey

        self.degen = def_degen if degen  is None else degen
        self.s02    = 1.0      if s02    is None else s02
        self.e0     = 0.0      if e0     is None else e0
        self.ei     = 0.0      if ei     is None else ei
        self.deltar = 0.0      if deltar is None else deltar
        self.sigma2 = 0.0      if sigma2 is None else sigma2
        self.third  = 0.0      if third  is None else third
        self.fourth = 0.0      if fourth is None else fourth

        self.k = None
        self.chi = None
        if self._feffdat is not None:
            self.create_spline_coefs()

    def __geom2label(self):
        """generate label by hashing path geometry"""
        rep = [self._feffdat.degen, self._feffdat.reff]
        for atom in self.geom:
            rep.extend(atom)

        for attr in ('s02', 'e0', 'ei', 'deltar', 'sigma2', 'third', 'fourth'):
            rep.append(getattr(self, attr, '_'))
        s = "|".join([str(i) for i in rep])
        return "p%s" % (b32hash(s)[:10].lower())

    def pathpar_name(self, parname):
        """
        get internal name of lmfit Parameter for a path paramter, using Path's hashkey
        """
        return f'{parname}_{self.hashkey}'

    def __copy__(self):
        return FeffPathGroup(filename=self.filename, label=self.label,
                             s02=self.s02, degen=self.degen, e0=self.e0,
                             ei=self.ei, deltar=self.deltar, sigma2=self.sigma2,
                             third=self.third, fourth=self.fourth)

    def __deepcopy__(self, memo):
        return FeffPathGroup(filename=self.filename, label=self.label,
                             s02=self.s02, degen=self.degen, e0=self.e0,
                             ei=self.ei, deltar=self.deltar, sigma2=self.sigma2,
                             third=self.third, fourth=self.fourth)

    @property
    def reff(self): return self._feffdat.reff

    @reff.setter
    def reff(self, val):  pass

    @property
    def nleg(self): return self._feffdat.nleg

    @nleg.setter
    def nleg(self, val):     pass

    @property
    def rmass(self): return self._feffdat.rmass

    @rmass.setter
    def rmass(self, val):  pass

    def __repr__(self):
        return f'<FeffPath Group label={self.label:s}, filename={self.filename:s}>'

    def create_path_params(self, params=None):
        """
        create Path Parameters within the current lmfit.Parameters namespace
        """
        if params is not None:
           self.params = params
        if self.params is None:
            self.params = Parameters()
        if self.params._asteval.symtable.get('sigma2_debye', None) is None:
            add_sigma2funcs(self.params)
        if self.label is None:
            self.label = self.__geom2label()
        self.store_feffdat()
        for pname in PATH_PARS:
            val =  getattr(self, pname)
            attr = 'value'
            if isinstance(val, str):
                attr = 'expr'
            kws =  {'vary': False, attr: val}
            parname = self.pathpar_name(pname)
            self.params.add(parname, **kws)
            self.params[parname].is_pathparam = True

    def create_spline_coefs(self):
        """pre-calculate spline coefficients for feff data"""
        self.spline_coefs = {}
        fdat = self._feffdat
        self.spline_coefs['pha'] = UnivariateSpline(fdat.k, fdat.pha, s=0)
        self.spline_coefs['amp'] = UnivariateSpline(fdat.k, fdat.amp, s=0)
        self.spline_coefs['rep'] = UnivariateSpline(fdat.k, fdat.rep, s=0)
        self.spline_coefs['lam'] = UnivariateSpline(fdat.k, fdat.lam, s=0)

    def store_feffdat(self):
        """stores data about this Feff path in the Parameters
        symbol table for use as `reff` and in sigma2 calcs
        """
        symtab = self.params._asteval.symtable
        symtab['feffpath'] = self._feffdat
        symtab['reff']  = self._feffdat.reff

    def __path_params(self, **kws):
        """evaluate path parameter value.  Returns
        (degen, s02, e0, ei, deltar, sigma2, third, fourth)
        """
        # put 'reff' and '_feffdat' into the symboltable so that
        # they can be used in constraint expressions
        self.store_feffdat()
        if self.params is None:
            self.create_path_params()
        out = []
        for pname in PATH_PARS:
            val = kws.get(pname, None)
            parname = self.pathpar_name(pname)
            if val is None:
                val = self.params[parname]._getval()
            out.append(val)
        return out

    def path_paramvals(self, **kws):
        (deg, s02, e0, ei, delr, ss2, c3, c4) = self.__path_params()
        return dict(degen=deg, s02=s02, e0=e0, ei=ei, deltar=delr,
                    sigma2=ss2, third=c3, fourth=c4)

    def report(self):
        "return  text report of parameters"
        tmpvals = self.__path_params()
        pathpars = {}
        for pname in ('degen', 's02', 'e0', 'deltar',
                      'sigma2', 'third', 'fourth', 'ei'):
            parname = self.pathpar_name(pname)
            if parname in self.params:
                pathpars[pname] = (self.params[parname].value, self.params[parname].stderr)

        out = [f" = Path '{self.label}' = ",
               f'    feffdat file = {self.filename}']
        geomlabel  = '    geometry  atom      x        y        z      ipot'
        geomformat = '            %4s      % .4f, % .4f, % .4f  %i'
        out.append(geomlabel)

        for atsym, iz, ipot, amass, x, y, z in self.geom:
            s = geomformat % (atsym, x, y, z, ipot)
            if ipot == 0: s = "%s (absorber)" % s
            out.append(s)

        stderrs = {}
        out.append('     {:7s}= {:s}'.format('reff',
                                              gformat(self._feffdat.reff)))

        for pname in ('degen', 's02', 'e0', 'r',
                      'deltar', 'sigma2', 'third', 'fourth', 'ei'):
            val = strval = getattr(self, pname, 0)
            parname = self.pathpar_name(pname)
            std = None
            if pname == 'r':
                parname = self.pathpar_name('deltar')
                par = self.params.get(parname, None)
                val = par.value + self._feffdat.reff
                strval = 'reff + ' + getattr(self, 'deltar', 0)
                std = par.stderr
            else:
                if pname in pathpars:
                    val, std = pathpars[pname]
                else:
                    par = self.params.get(parname, None)
                    if par is not None:
                        val = par.value
                        std = par.stderr

            if std is None  or std <= 0:
                svalue = gformat(val)
            else:
                svalue = "{:s} +/-{:s}".format(gformat(val), gformat(std))
            if pname == 's02':
                pname = 'n*s02'

            svalue = "     {:7s}= {:s}".format(pname, svalue)
            if isinstance(strval, str):
                svalue = "{:s}  := '{:s}'".format(svalue, strval)

            if val == 0 and pname in ('third', 'fourth', 'ei'):
                continue
            out.append(svalue)
        return '\n'.join(out)

    def calc_chi_from_params(self, params, **kws):
        "calculate chi(k) from Parameters, ParameterGroup, and/or kws for path parameters"
        if isinstance(params, Parameters):
            self.create_path_params(params=params)
        else:
            self.create_path_params(params=group2params(params))
        self._calc_chi(**kws)

    def _calc_chi(self, k=None, kmax=None, kstep=None, degen=None, s02=None,
                 e0=None, ei=None, deltar=None, sigma2=None,
                 third=None, fourth=None, debug=False, interp='cubic', **kws):
        """calculate chi(k) with the provided parameters"""
        fdat = self._feffdat
        if fdat.reff < 0.05:
            print('reff is too small to calculate chi(k)')
            return
        # make sure we have a k array
        if k is None:
            if kmax is None:
                kmax = 30.0
            kmax = min(max(fdat.k), kmax)
            if kstep is None: kstep = 0.05
            k = kstep * np.arange(int(1.01 + kmax/kstep), dtype='float64')

        reff = fdat.reff
        # get values for all the path parameters
        (degen, s02, e0, ei, deltar, sigma2, third, fourth)  = \
                self.__path_params(degen=degen, s02=s02, e0=e0, ei=ei,
                                 deltar=deltar, sigma2=sigma2,
                                 third=third, fourth=fourth)

        # create e0-shifted energy and k, careful to look for |e0| ~= 0.
        en = k*k - e0*ETOK
        if min(abs(en)) < SMALL_ENERGY:
            try:
                en[np.where(abs(en) < 1.5*SMALL_ENERGY)] = SMALL_ENERGY
            except ValueError:
                pass
        # q is the e0-shifted wavenumber
        q = np.sign(en)*np.sqrt(abs(en))

        # lookup Feff.dat values (pha, amp, rep, lam)
        if interp.startswith('lin'):
            pha = np.interp(q, fdat.k, fdat.pha)
            amp = np.interp(q, fdat.k, fdat.amp)
            rep = np.interp(q, fdat.k, fdat.rep)
            lam = np.interp(q, fdat.k, fdat.lam)
        else:
            pha = self.spline_coefs['pha'](q)
            amp = self.spline_coefs['amp'](q)
            rep = self.spline_coefs['rep'](q)
            lam = self.spline_coefs['lam'](q)

        if debug:
            self.debug_k   = q
            self.debug_pha = pha
            self.debug_amp = amp
            self.debug_rep = rep
            self.debug_lam = lam

        # p = complex wavenumber, and its square:
        pp   = (rep + 1j/lam)**2 + 1j * ei * ETOK
        p    = np.sqrt(pp)

        # the xafs equation:
        cchi = np.exp(-2*reff*p.imag - 2*pp*(sigma2 - pp*fourth/3) +
                      1j*(2*q*reff + pha +
                          2*p*(deltar - 2*sigma2/reff - 2*pp*third/3) ))

        cchi = degen * s02 * amp * cchi / (q*(reff + deltar)**2)
        cchi[0] = 2*cchi[1] - cchi[2]
        # outputs:
        self.k = k
        self.p = p
        self.chi = cchi.imag
        self.chi_imag = -cchi.real

def path2chi(path, paramgroup=None, **kws):
    """calculate chi(k) for a Feff Path,
    optionally setting path parameter values
    output chi array will be written to path group

    Parameters:
    ------------
      path:        a FeffPath Group
      params:      lmfit Parameters or larch ParameterGroup
      kmax:        maximum k value for chi calculation [20].
      kstep:       step in k value for chi calculation [0.05].
      k:           explicit array of k values to calculate chi.

    Returns:
    ---------
      None - outputs are written to path group

    """
    if not isNamedClass(path, FeffPathGroup):
        msg('%s is not a valid Feff Path' % path)
        return
    path.calc_chi_from_params(paramgroup, **kws)

def ff2chi(paths, group=None, paramgroup=None, k=None, kmax=None,
            kstep=0.05, _larch=None, **kws):
    """sum chi(k) for a list of FeffPath Groups.

    Parameters:
    ------------
      paths:       a list of FeffPath Groups or dict of {label: FeffPathGroups}
      paramgroup:  a Parameter Group for calculating Path Parameters [None]
      kmax:        maximum k value for chi calculation [20].
      kstep:       step in k value for chi calculation [0.05].
      k:           explicit array of k values to calculate chi.
    Returns:
    ---------
       group contain arrays for k and chi

    This essentially calls path2chi() for each of the paths in the
    `paths` and writes the resulting arrays to group.k and group.chi.

    """
    params = group2params(paramgroup)

    if isinstance(paths, (list, tuple)):
        pathlist = paths
    elif isinstance(paths, dict):
        pathlist = list(paths.values())
    else:
        raise ValueErrror('paths must be list, tuple, or dict')

    for path in pathlist:
        if not isNamedClass(path, FeffPathGroup):
            print('%s is not a valid Feff Path' % path)
            return
        path.create_path_params(params=params)
        path._calc_chi(k=k, kstep=kstep, kmax=kmax)
    k = pathlist[0].k[:]
    out = np.zeros_like(k)
    for path in pathlist:
        out += path.chi

    if group is None:
        group = Group()
    else:
        group = set_xafsGroup(group, _larch=_larch)
    group.k = k
    group.chi = out
    return group

def feffpath(filename=None, label=None, s02=None, degen=None,
             e0=None,ei=None, deltar=None, sigma2=None, third=None,
             fourth=None, _larch=None, **kws):
    """create a Feff Path Group from a *feffNNNN.dat* file.

    Parameters:
    -----------
      filename:  name (full path of) *feffNNNN.dat* file
      label:     label for path   [file name]
      degen:     path degeneracy, N [taken from file]
      s02:       S_0^2    value or parameter [1.0]
      e0:        E_0      value or parameter [0.0]
      deltar:    delta_R  value or parameter [0.0]
      sigma2:    sigma^2  value or parameter [0.0]
      third:     c_3      value or parameter [0.0]
      fourth:    c_4      value or parameter [0.0]
      ei:        E_i      value or parameter [0.0]

    For all the options described as **value or parameter** either a
    numerical value or a Parameter (as created by param()) can be given.

    Returns:
    ---------
        a FeffPath Group.

    """
    return FeffPathGroup(filename=filename, label=label, s02=s02,
                         degen=degen, e0=e0, ei=ei, deltar=deltar,
                         sigma2=sigma2, third=third, fourth=fourth)


import numpy as np
import tensorflow as tf
from parameters import par
from itertools import product

class AdamOpt:

    """
    Example of use:

    optimizer = AdamOpt.AdamOpt(variables, learning_rate=self.lr)
    self.train = optimizer.compute_gradients(self.loss, gate=0)
    gvs = optimizer.return_gradients()
    self.g = gvs[0][0]
    self.v = gvs[0][1]
    """

    def __init__(self, variables, learning_rate = 0.001):

        self.beta1 = 0.9
        self.beta2 = 0.999
        self.epsilon = 1e-08
        self.t = 0
        self.variables = variables
        self.learning_rate = learning_rate

        self.m = {}
        self.v = {}
        self.delta_grads = {}
        for var in self.variables:
            self.m[var.op.name]  = tf.Variable(tf.zeros(var.get_shape()), trainable=False)
            self.v[var.op.name]  = tf.Variable(tf.zeros(var.get_shape()), trainable=False)
            self.delta_grads[var.op.name]  = tf.Variable(tf.zeros(var.get_shape()), trainable=False)

        self.grad_descent = tf.train.GradientDescentOptimizer(learning_rate = 1.0)

    def reset_params(self):

        self.t = 0
        reset_op = []
        for var in self.variables:
            reset_op.append(tf.assign(self.m[var.op.name], tf.zeros(var.get_shape())))
            reset_op.append(tf.assign(self.v[var.op.name], tf.zeros(var.get_shape())))
            reset_op.append(tf.assign(self.delta_grads[var.op.name], tf.zeros(var.get_shape())))

        return tf.group(*reset_op)

    def optimize(self, loss):

        grads_and_vars = self.compute_gradients(loss)
        train_op = self.apply_gradients(grads_and_vars)

        return train_op

    def compute_gradients(self, loss):

        self.gradients = self.grad_descent.compute_gradients(loss, var_list = self.variables)

        self.t += 1
        lr = self.learning_rate*np.sqrt(1-self.beta2**self.t)/(1-self.beta1**self.t)
        self.update_var_op = []

        #grads_and_vars = []
        for (grads, _), var in zip(self.gradients, self.variables):
            new_m = self.beta1*self.m[var.op.name] + (1-self.beta1)*grads
            new_v = self.beta2*self.v[var.op.name] + (1-self.beta2)*grads*grads

            delta_grad = - lr*new_m/(tf.sqrt(new_v) + self.epsilon)
            delta_grad = tf.clip_by_norm(delta_grad, 1)

            self.update_var_op.append(tf.assign(self.m[var.op.name], new_m))
            self.update_var_op.append(tf.assign(self.v[var.op.name], new_v))

            if 'W_rnn' in var.op.name:
                print('Applied W_rnn mask.')
                delta_grad *= par['W_rnn_mask']
            elif 'W_in' in var.op.name:
                print('Applied W_in mask.')
                delta_grad *= par['W_in_mask']
            elif 'W_d_rnn' in var.op.name:
                print('Applied W_d_rnn mask.')
                delta_grad *= par['W_d_rnn_mask']
            elif 'W_out' in var.op.name:
                print('Applied W_out mask.')
                delta_grad *= par['W_out_mask']

            self.update_var_op.append(tf.assign(self.delta_grads[var.op.name], delta_grad))
            self.update_var_op.append(tf.assign_add(var, delta_grad))

        return tf.group(*self.update_var_op)

    def apply_gradients(self, grads_and_vars):
        # currently not in use
        for (grad, var) in grads_and_vars:
            if 'W_rnn' in var.op.name:
                print('Applied W_rnn mask.')
                grad *= par['W_rnn_mask']
            elif 'W_in' in var.op.name:
                print('Applied W_in mask.')
                grad *= par['W_in_mask']
            elif 'W_out' in var.op.name:
                print('Applied W_out mask.')
                grad *= par['W_out_mask']
            self.update_var_op.append(tf.assign_add(var, grad))

        return tf.group(*self.update_var_op)

    def return_delta_grads(self):
        return self.delta_grads

    def return_means(self):
        return self.m

    def return_grads_and_vars(self):
        return self.gradients