# Copyright (c) 2020 Matthew Earl
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
#     The above copyright notice and this permission notice shall be included
#     in all copies or substantial portions of the Software.
#
#     THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
#     OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
#     MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
#     NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
#     DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
#     OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
#     USE OR OTHER DEALINGS IN THE SOFTWARE.

import functools
from typing import NamedTuple, Iterable

import numpy as np

from . import bsp

class Ray(NamedTuple):
    origin: np.ndarray
    dir: np.ndarray

def _infront(point, plane_norm, plane_dist):
    return np.dot(point, plane_norm) - plane_dist >= 0

class _GlowTexture(NamedTuple):
    im: np.ndarray
    sample_distribution: np.ndarray
    emissivity: float

    @classmethod
    def from_im(cls, im):
        dist = np.sum(im, axis=2)
        em = np.sum(dist)
        dist /= em
        return cls(im, dist, em)

class _FaceTextureCache:
    def __init__(self, pal: np.ndarray):
        self._pal = pal
        self._glow_pal = np.copy(pal)
        self._glow_pal[:-32, :] = 0.
        self._cache = {}

    def _load_face_texture(self, face: bsp.Face, glow: bool = False):
        texture = face.tex_info.texture
        indices = np.fromstring(texture.data[0], dtype=np.uint8).reshape((texture.height, texture.width))
        if glow and np.all(indices < 255 - 32):
            return None
        pal = self._glow_pal if glow else self._pal
        array_im = pal[np.fromstring(texture.data[0], dtype=np.uint8).reshape((texture.height, texture.width))]
        #array_im = array_im ** 0.8
        return _GlowTexture.from_im(array_im) if glow else array_im

    def get_face_texture(self, face: bsp.Face, *, glow: bool = False):
        key = face, glow
        if key not in self._cache:
            self._cache[key] = self._load_face_texture(face, glow)
        return self._cache[key]

def _sample_texture(texture_cache: _FaceTextureCache, face: bsp.Face, poi: np.ndarray, glow=False):
    s, t = face.tex_info.vert_to_tex_coords(poi)
    if glow:
        im = texture_cache.get_face_texture(face, glow=True).im
    else:
        im = texture_cache.get_face_texture(face)
    s %= im.shape[1]
    t %= im.shape[0]
    return im[int(t), int(s)]

@functools.lru_cache(None)
def _get_light_faces_and_textures(texture_cache: _FaceTextureCache, leaf: bsp.Leaf):
    faces = leaf.visible_faces
    textures = [texture_cache.get_face_texture(face, glow=True) for face in faces]
    faces = [f for f, t in zip(faces, textures) if t is not None]
    textures = [t for t in textures if t is not None]
    return faces, textures
    return faces[12:13], textures[12:13]

def _sample_light_sources2(texture_cache: _FaceTextureCache, leaf: bsp.Leaf, normal: np.ndarray, poi: np.ndarray):
    #import ipdb; ipdb.set_trace()
    faces, textures = _get_light_faces_and_textures(texture_cache, leaf)

    # First choose a face to sample from
    em = np.array([t.emissivity for t in textures])
    ray_dir = np.array([face.centroid for face in faces]) - poi
    face_normals = np.stack([face.plane[0] for face in faces])
    distance = np.linalg.norm(ray_dir, axis=1)
    ray_dir /= distance[:, None]
    light_cos_angle = np.abs(np.sum(ray_dir * face_normals, axis=1))
    source_cos_angle = np.abs(np.sum(ray_dir * normal, axis=1))
    area = np.array([face.area for face in faces])

    face_probs = area * em * light_cos_angle * source_cos_angle / distance ** 2
    face_probs /= np.sum(face_probs)

    idx = np.random.choice(range(len(faces)), p=face_probs)
    light_face = faces[idx]
    light_texture = textures[idx]
    face_prob = face_probs[idx]

    # Then pick a random point on the face
    for i in range(20):
        vert_array = np.array(list(light_face.vertices))
        p = np.random.random((len(vert_array),))
        p /= np.sum(p)
        light_pos = p @ vert_array
        colour = _sample_texture(texture_cache, light_face, light_pos, glow=True)
        if not np.array_equal(colour, (0, 0, 0)):
            break

    return light_face, light_pos, colour, (i + 1) * face_prob / light_face.area

def _sample_light_sources(texture_cache: _FaceTextureCache, leaf: bsp.Leaf, normal: np.ndarray, poi: np.ndarray):
    faces = leaf.visible_faces
    textures = [texture_cache.get_face_texture(face, glow=True) for face in faces]
    faces = [f for f, t in zip(faces, textures) if t is not None]
    textures = [t for t in textures if t is not None]

    # First choose a face to sample from
    em = np.array([t.emissivity for t in textures])
    ray_dir = np.array([face.centroid for face in faces]) - poi
    face_normals = np.stack([face.plane[0] for face in faces])
    distance = np.linalg.norm(ray_dir, axis=1)
    ray_dir /= distance[:, None]
    light_cos_angle = np.maximum(np.sum(ray_dir * face_normals, axis=1), 0)
    source_cos_angle = np.maximum(-np.sum(ray_dir * normal, axis=1), 0)
    area = np.array([face.area for face in faces])

    #print(f"area {area}\n em {em}\n light_cos_angle {light_cos_angle}\n source_cos_angle {source_cos_angle}\n distance {distance}")
    face_probs = area * em * light_cos_angle * source_cos_angle / distance ** 2
    face_probs /= np.sum(face_probs)

    idx = np.random.choice(range(len(faces)), p=face_probs)
    light_face = faces[idx]
    light_texture = textures[idx]
    face_prob = face_probs[idx]

    # Sample a pixel on the face
    reshaped_im = light_texture.im.reshape((-1, 3))
    pixel_probs = light_texture.sample_distribution.ravel()
    idx = np.random.choice(range(reshaped_im.shape[0]), p=pixel_probs)
    colour = reshaped_im[idx]
    pixel_prob = pixel_probs[idx]

    # Work out the position
    s, t = idx // light_texture.im.shape[1], idx % light_texture.im.shape[1]
    light_pos = light_face.tex_info.tex_coords_to_vert((s, t))

    return light_face, light_pos, colour, face_prob * pixel_prob / light_face.tex_info.texel_area

def _sample_direct_light_contribution(model, texture_cache, face, poi, leaf):
    light_face, light_pos, light_colour, prob = _sample_light_sources2(texture_cache, leaf, face.plane[0], poi)

    ray_dir = light_pos - poi
    ray_len = np.linalg.norm(ray_dir)
    ray_dir /= ray_len
    ray = Ray(poi, ray_dir)
    trace_face, _, _, _ = _ray_bsp_intersect(model, ray, 1e-3)

    if trace_face != light_face:
        out = np.zeros((3,))
    else:
        scale = np.abs(np.dot(ray_dir, light_face.plane[0]) * np.dot(ray_dir, face.plane[0]))
        scale /= np.linalg.norm(light_pos - poi) ** 2
        scale /= prob
        out = 10 * light_colour * scale

    return out

def _trace_leaves(is_leaf: bool, node: bsp.Node, ray: Ray, near_clip: float = 0., far_clip: float = np.inf):
    """Trace a ray through a BSP node, yielding all encountered leaves (in near to far order)."""
    if is_leaf:
        yield node, near_clip, far_clip
    else:
        plane = node.plane
        plane_norm = np.array(plane.normal)
        plane_dist = plane.dist

        beta = np.dot(plane_norm, ray.dir)
        near_child, far_child = (1, 0) if beta > 0 else (0, 1)

        if np.abs(beta) < 1e-5:
            child = 0 if _infront(ray.origin, plane_norm, plane_dist) else 1
            yield from _trace_leaves(node.child_is_leaf(child), node.get_child(child), ray, near_clip, far_clip)
        else:
            alpha = (plane_dist - np.dot(plane_norm, ray.origin)) / beta

            if near_clip <= alpha < far_clip:
                yield from _trace_leaves(node.child_is_leaf(near_child), node.get_child(near_child),
                                         ray, near_clip, alpha)
                yield from _trace_leaves(node.child_is_leaf(far_child), node.get_child(far_child),
                                         ray, alpha, far_clip)
            elif alpha < near_clip:
                yield from _trace_leaves(node.child_is_leaf(far_child), node.get_child(far_child),
                                         ray, near_clip, far_clip)
            else:
                yield from _trace_leaves(node.child_is_leaf(near_child), node.get_child(near_child),
                                         ray, near_clip, far_clip)

def _ray_face_intersect(face: bsp.Face, ray: Ray):
    face_norm, face_dist = face.plane

    beta = np.dot(face_norm, ray.dir)
    if abs(beta) >= 1e-5:
        ray_dist = (face_dist - np.dot(face_norm, ray.origin)) / beta
        if ray_dist >= 0:
            poi = ray.origin + ray.dir * ray_dist
            if all(_infront(poi, n, d) for n, d in face.edge_planes):
                return poi, ray_dist

    return None, np.inf

def _ray_faces_intersect(faces: Iterable[bsp.Face], ray: Ray, near_clip: float = 0., far_clip: float = np.inf):
    nearest_face, nearest_poi, nearest_dist = None, None, np.inf
    for face in faces:
        poi, ray_dist = _ray_face_intersect(face, ray)
        if ray_dist < nearest_dist and near_clip - 1e-4 <= ray_dist <= far_clip + 1e-4:
            nearest_face, nearest_poi, nearest_dist = face, poi, ray_dist
    return nearest_face, nearest_poi, nearest_dist

def _ray_bsp_intersect(model: bsp.Model, ray: Ray, near_clip=0., far_clip=np.inf):
    nearest_leaf = None
    for leaf, near_clip, far_clip in _trace_leaves(False, model.node, ray, near_clip, far_clip):
        nearest_face, nearest_poi, nearest_dist = _ray_faces_intersect(leaf.faces, ray, near_clip, far_clip)
        if nearest_face is not None:
            nearest_leaf = leaf
            break

    return nearest_face, nearest_poi, nearest_dist, nearest_leaf

def raytracer_main2():
    import io
    import sys
    import logging

    import cv2

    from .bsp import Bsp
    from . import pak

    root_logger = logging.getLogger()
    root_logger.addHandler(logging.StreamHandler())
    root_logger.setLevel(logging.DEBUG)

    fs = pak.Filesystem(sys.argv[1])
    bsp = Bsp(io.BytesIO(fs[sys.argv[2]]))

    WIDTH, HEIGHT = 200, 200
    K_inv = np.matrix([[WIDTH / 2.,             0,  WIDTH  / 2],
                       [0,           HEIGHT / 2.,  HEIGHT / 2],
                       [0,                      0,  1]]).I
    K_inv = np.array(K_inv)

    player_start = next(iter(e for e in bsp.entities if e['classname'] == 'info_player_start'))
    ray_origin = np.array(player_start['origin']) + [0, 0, 21]

    rot = np.array([[1., 0.,  0.],
                    [0., 0.,  1.],
                    [0., -1., 0.]])

    pal = np.fromstring(fs['gfx/palette.lmp'], dtype=np.uint8).reshape(256, 3) / 255
    texture_cache = _FaceTextureCache(pal)

    def pix_to_dir(x, y):
        ray_dir = rot @ K_inv @ np.array([x, y, 1])
        return ray_dir / np.linalg.norm(ray_dir)

    origin_leaf = bsp.models[0].get_leaf_from_point(ray_origin)
    visible_leaves = list(origin_leaf.visible_leaves)
    vis_faces = list({f for l in origin_leaf.visible_leaves for f in l.faces})
    print("Num visible leaves %s %s, num vis faces", len(visible_leaves), len(set(visible_leaves)), len(vis_faces))

    cv2.namedWindow("out")
    out = np.zeros((HEIGHT, WIDTH, 3))
    color_wheel = np.random.random((32, 3))
    color_wheel /= np.max(color_wheel, axis=1)[:, None]
    for y in range(HEIGHT):
        print(y)
        for x in range(WIDTH):
            #x, y = 150, 120
            ray = Ray(ray_origin, pix_to_dir(x, y))

            face, poi, dist, leaf = _ray_bsp_intersect(bsp.models[0], ray)
            #face, poi, dist = _ray_faces_intersect(vis_faces, ray)

            #color = (0., 0., 0.) if face is None else color_wheel[hash(face) % len(color_wheel)]
            #color = (1, 0, 1) if face is None else _sample_texture(texture_cache, face, poi)
            color = (1, 0, 1) if face is None else (1, 1, 1)

            if face is not None:
                # Sample light sources
                color = np.array([_sample_direct_light_contribution(
                                    bsp.models[0], texture_cache, face, poi, leaf)
                                    for i in range(1)]).mean(axis=0)
            out[y, x] = np.flip(color, axis=0)

        print(np.max(out[y]))
        cv2.imshow("out", out)
        cv2.waitKey(1)

    cv2.destroyWindow("out")
    cv2.imwrite("out.png", out * 255.)

def raytracer_main():
    profiler = 0
    if profiler == 0:
        raytracer_main2()
    elif profiler == 1:
        from line_profiler import LineProfiler
        profile = LineProfiler(_sample_light_sources2,
                               _FaceTextureCache.get_face_texture)
        profile.runctx('raytracer_main2()', globals(), locals())
        profile.print_stats()
    else:
        import cProfile
        cProfile.runctx('raytracer_main2()', globals(), locals(), 'stats5')


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

from __future__ import print_function
import mbuild as mb
import builtins as __builtin__
import numpy as np
from foyer import Forcefield
#All the old import from compound
import collections
from collections import OrderedDict, defaultdict
from copy import deepcopy
import itertools
import os
import sys
import tempfile
from warnings import warn
##
from mbuild.utils.io import run_from_ipython, import_

r0BB_pvpy = 1.0
r0BB_pvph = 1.134
r0BBHB_pvpy = 0.37
r0BBHB_pvph = 0.42

# Define backbone (BB) classes
class BB(mb.Compound):
    def __init__(self):
        super(BB, self).__init__(pos=[0,0,0],name = 'bb')#Initizlize an instance of abstract class

        bead = mb.Particle(pos=[0.0, 0.0, 0.0], name='BB')
        self.add(bead)

        port = mb.Port(anchor=list(self.particles_by_name('BB'))[0], orientation=[0, 1, 0], separation=r0BB/2)
        self.add(port, 'up')
        port = mb.Port(anchor=list(self.particles_by_name('BB'))[0], orientation=[1, 0, 0], separation=r0BBHB/2)
        self.add(port,'toHB')
        port = mb.Port(anchor=list(self.particles_by_name('BB'))[0], orientation=[0, -1, 0], separation=r0BB/2)
        self.add(port, 'down')

class BB_pvpy(BB):
    def __init__(self):
        super(BB_pvpy,self).__init__(r0BB=r0BB_pvpy)
        for par in self.particles():
            par.name = '_BB_pvpy'
            #print(par.name)

        # Name the entire compound (particle+its ports)
        self.name = '_bb_pvpy'

class BB_pvpy():
    def __init__(self):
        super(BB_pvpy, self).__init__(pos=[0,0,0],name = 'bb_pvpy')#Initizlize an instance of abstract class

        bead = mb.Particle(pos=[0.0, 0.0, 0.0], name='BB_pvpy')
        self.add(bead)

        port = mb.Port(anchor=list(self.particles_by_name('BB_pvpy'))[0], orientation=[0, 1, 0], separation=r0BB_pvpy/2)
        self.add(port, 'up')
        port = mb.Port(anchor=list(self.particles_by_name('BB_pvpy'))[0], orientation=[1, 0, 0], separation=r0BBHB_pvpy/2)
        self.add(port,'toHB')
        port = mb.Port(anchor=list(self.particles_by_name('BB_pvpy'))[0], orientation=[0, -1, 0], separation=r0BB_pvpy/2)
        self.add(port, 'down')

class BB_pvph(mb.Compound):
    def __init__(self):
        super(BB_pvph, self).__init__(pos=[0,0,0],name = 'bb_pvph')#Initizlize an instance of abstract class

        bead = mb.Particle(pos=[0.0, 0.0, 0.0], name='BB_pvph')
        self.add(bead)

        port = mb.Port(anchor=list(self.particles_by_name('BB_pvph'))[0], orientation=[0, 1, 0], separation=r0BB_pvph/2)
        self.add(port, 'up')
        port = mb.Port(anchor=list(self.particles_by_name('BB_pvph'))[0], orientation=[1, 0, 0], separation=r0BBHB_pvph/2)
        self.add(port,'toHB')
        port = mb.Port(anchor=list(self.particles_by_name('BB_pvph'))[0], orientation=[0, -1, 0], separation=r0BB_pvph/2)
        self.add(port, 'down')

# Define hydrogen bonding (HB) classes
class HB_pvpy(mb.Compound):
    def __init__(self):
        super(HB_pvpy, self).__init__(pos=[0.0, 0.0, 0.0], name='hb_pvpy')

        bead = mb.Particle(pos=[0.0, 0.0, 0.0], name='HB_pvpy')
        self.add(bead)

        port = mb.Port(anchor=list(self.particles_by_name('HB_pvpy'))[0], orientation=[-1, 0, 0], separation=r0BBHB_pvpy/2)
        self.add(port, 'toBB')

class HB_pvph(mb.Compound):
    def __init__(self):
        super(HB_pvph, self).__init__(pos=[0.0, 0.0, 0.0], name='hb_pvph')

        bead = mb.Particle(pos=[0.0, 0.0, 0.0], name='HB_pvph')
        self.add(bead)

        port = mb.Port(anchor=list(self.particles_by_name('HB_pvph'))[0], orientation=[-1, 0, 0], separation=r0BBHB_pvph/2)
        self.add(port, 'toBB')

# Define monomer (M) classes - combine backbone and hydrogen bonding beads
class M_pvpy(mb.Compound):
    def __init__(self):
        super(M_pvpy,self).__init__()
        bb = BB_pvpy()
        hb = HB_pvpy()
        self.add((bb,hb))
        #Move
        mb.force_overlap(move_this=hb, from_positions=hb['toBB'],to_positions=bb['toHB'])

class M_pvph(mb.Compound):
    def __init__(self):
        super(M_pvph,self).__init__()
        bb = BB_pvph()
        hb = HB_pvph()
        self.add((bb,hb))
        #Move
        mb.force_overlap(move_this=hb, from_positions=hb['toBB'],to_positions=bb['toHB'])

class polymer_pvpy(mb.Compound):
    chain_len = None
    def __init__(self,length=None):
        self.chain_len = length
        super(polymer_pvpy,self).__init__()
        if length != 0:
            last_M = M_pvpy()
            self.add(last_M)

            for i in range(length-1):
                new_M = M_pvpy()
                self.add(new_M)
                mb.force_overlap(move_this = new_M, from_positions=(new_M.all_ports())[0],to_positions = (last_M.all_ports())[-1])
                last_M = new_M

    def polymer_visualize_py3dmol(self, show_ports=False):
        """Visualize the Compound using py3Dmol.
        Allows for visualization of a Compound within a Jupyter Notebook.
        Parameters
        ----------
        show_ports : bool, optional, default=False
            Visualize Ports in addition to Particles
        color_scheme : dict, optional
            Specify coloring for non-elemental particles
            keys are strings of the particle names
            values are strings of the colors
            i.e. {'_CGBEAD': 'blue'}
        Returns
        ------
        view : py3Dmol.view
        """
        py3Dmol = import_('py3Dmol')
        remove_digits = lambda x: ''.join(i for i in x if not i.isdigit()
                                              or i == '_')

        for particle in self.particles():
            if not particle.name:
                particle.name = 'UNK'

        view = py3Dmol.view()
        rad = {'BB_pvpy':.5,'HB_pvpy':0.22}
        col = {'BB_pvpy':'#ffff14','HB_pvpy':'#000000'}

        for p in self.particles(include_ports=False):
            view.addSphere({
                'center': {'x':p.pos[0], 'y':p.pos[1], 'z':p.pos[2]},
                'radius' :rad[p.name],
                'color': col[p.name],
                'alpha': 0.9})
        view.zoomTo()
        view.show()

        return view

class polymer_pvph(mb.Compound):
    chain_len = None
    def __init__(self,length=None):
        self.chain_len = length
        super(polymer_pvph,self).__init__()
        if length != 0:
            last_M = M_pvph()
            self.add(last_M)

            for i in range(length-1):
                new_M = M_pvph()
                self.add(new_M)
                mb.force_overlap(move_this = new_M, from_positions=(new_M.all_ports())[0],to_positions = (last_M.all_ports())[-1])
                last_M = new_M

    def polymer_visualize_py3dmol(self, show_ports=False):
        """Visualize the Compound using py3Dmol.
        Allows for visualization of a Compound within a Jupyter Notebook.
        Parameters
        ----------
        show_ports : bool, optional, default=False
            Visualize Ports in addition to Particles
        color_scheme : dict, optional
            Specify coloring for non-elemental particles
            keys are strings of the particle names
            values are strings of the colors
            i.e. {'_CGBEAD': 'blue'}
        Returns
        ------
        view : py3Dmol.view
        """
        py3Dmol = import_('py3Dmol')
        remove_digits = lambda x: ''.join(i for i in x if not i.isdigit()
                                              or i == '_')

        for particle in self.particles():
            if not particle.name:
                particle.name = 'UNK'

        view = py3Dmol.view()
        rad = {'BB_pvph':.567,'HB_pvph':0.24}
        col = {'BB_pvph':'#15b01a','HB_pvph':'#e50000'}

        for p in self.particles(include_ports=False):
            view.addSphere({
                'center': {'x':p.pos[0], 'y':p.pos[1], 'z':p.pos[2]},
                'radius' :rad[p.name],
                'color': col[p.name],
                'alpha': 0.9})
        view.zoomTo()
        view.show()

        return view

class polymer_box(mb.Compound):
    def __init__(self,chain_len_pvpy = [],chain_len_pvph = [],dim = [0,0,0]):
        super(polymer_box,self).__init__()
        self.chain_len_pvpy = chain_len_pvpy
        self.chain_len_pvph = chain_len_pvph
        if (len(chain_len_pvpy) + len(chain_len_pvph)) > 0 and dim[0]*dim[1]*dim[2] != (len(chain_len_pvpy) + len(chain_len_pvph)):
            dim = [len(chain_len_pvpy) + len(chain_len_pvph),1,1]
        chain_num_pvpy = 0
        chain_num_pvph = 0
        total_chain_num = 0
        for i in range(dim[0]):
            for j in range(dim[1]):
                for k in range(dim[2]):
                    if (len(chain_len_pvpy) > 0) | (len(chain_len_pvph) > 0):
                        if (total_chain_num < len(chain_len_pvpy)):
                            new_polymer = polymer_pvpy(chain_len_pvpy[chain_num_pvpy])
                            new_polymer.translate_to([i*3.,j*3.,k*3.])
                            self.add(new_polymer)
                            chain_num_pvpy += 1
                            total_chain_num += 1
                        elif ((total_chain_num - len(chain_len_pvpy)) < len(chain_len_pvph)):
                            new_polymer = polymer_pvph(chain_len_pvph[chain_num_pvph])
                            new_polymer.translate_to([i*3.,j*3.,k*3.])
                            self.add(new_polymer)
                            chain_num_pvph += 1
                            total_chain_num += 1

    def visualize(self):
        py3Dmol = import_('py3Dmol')
        view = py3Dmol.view()
        if (len(self.chain_len_pvpy) > 0) and (len(self.chain_len_pvph) > 0):
            rad = {'BB_pvpy':.5,'BB_pvph':.567,'HB_pvpy':0.22,'HB_pvph':0.24}
            col = {'BB_pvph':'#15b01a','BB_pvpy':'#ffff14','HB_pvph':'#e50000','HB_pvpy':'#000000'}
        elif (len(self.chain_len_pvpy) == 0):
            rad = {'BB_pvph':.567,'HB_pvph':0.24}
            col = {'BB_pvph':'#15b01a','HB_pvph':'#e50000'}
        elif (len(self.chain_len_pvph) == 0):
            rad = {'BB_pvpy':.5,'HB_pvpy':0.22}
            col = {'BB_pvpy':'#ffff14','HB_pvpy':'#000000'}

        remove_digits = lambda x: ''.join(i for i in x if not i.isdigit()
                                              or i == '_')

        #modified_color_scheme = {}

        for chain in self.children:
            for particle in chain.particles():
                #particle.name = remove_digits(particle.name).upper()
                if not particle.name:
                    particle.name = 'UNK'

                for p in chain.particles(include_ports=False):
                    view.addSphere({
                        'center': {'x':p.pos[0], 'y':p.pos[1], 'z':p.pos[2]},
                        'radius' :rad[p.name],
                        'color': col[p.name],
                        'alpha': 0.9})
        view.zoomTo()
        view.show()
        return view


# *****************************************************************************
#  Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of the NVIDIA CORPORATION nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# *****************************************************************************

import functools
import json
import re
from pathlib import Path

import librosa
import numpy as np
import torch
import torch.nn.functional as F
from scipy import ndimage
from scipy.stats import betabinom

import common.layers as layers
from common.text.text_processing import TextProcessing
from common.utils import load_wav_to_torch, load_filepaths_and_text, to_gpu
from sklearn.preprocessing import StandardScaler
import numpy as np

class BetaBinomialInterpolator:
    """Interpolates alignment prior matrices to save computation.

    Calculating beta-binomial priors is costly. Instead cache popular sizes
    and use img interpolation to get priors faster.
    """
    def __init__(self, round_mel_len_to=100, round_text_len_to=20):
        self.round_mel_len_to = round_mel_len_to
        self.round_text_len_to = round_text_len_to
        self.bank = functools.lru_cache(beta_binomial_prior_distribution)

    def round(self, val, to):
        return max(1, int(np.round((val + 1) / to))) * to

    def __call__(self, w, h):
        bw = self.round(w, to=self.round_mel_len_to)
        bh = self.round(h, to=self.round_text_len_to)
        ret = ndimage.zoom(self.bank(bw, bh).T, zoom=(w / bw, h / bh), order=1)
        assert ret.shape[0] == w, ret.shape
        assert ret.shape[1] == h, ret.shape
        return ret

def beta_binomial_prior_distribution(phoneme_count, mel_count, scaling=1.0):
    P = phoneme_count
    M = mel_count
    x = np.arange(0, P)
    mel_text_probs = []
    for i in range(1, M+1):
        a, b = scaling * i, scaling * (M + 1 - i)
        rv = betabinom(P, a, b)
        mel_i_prob = rv.pmf(x)
        mel_text_probs.append(mel_i_prob)
    return torch.tensor(np.array(mel_text_probs))

def estimate_pitch(wav, mel_len, method='pyin', normalize_mean=None,
                   normalize_std=None, n_formants=1):

    if type(normalize_mean) is float or type(normalize_mean) is list:
        normalize_mean = torch.tensor(normalize_mean)

    if type(normalize_std) is float or type(normalize_std) is list:
        normalize_std = torch.tensor(normalize_std)

    if method == 'pyin':

        snd, sr = librosa.load(wav)
        pitch_mel, voiced_flag, voiced_probs = librosa.pyin(
            snd, fmin=90,
            fmax=600, frame_length=1024)
        assert np.abs(mel_len - pitch_mel.shape[0]) <= 1.0

        pitch_mel = np.where(np.isnan(pitch_mel), 0.0, pitch_mel)
        pitch_mel = torch.from_numpy(pitch_mel).unsqueeze(0)
        pitch_mel = F.pad(pitch_mel, (0, mel_len - pitch_mel.size(1)))

        if n_formants > 1:
            raise NotImplementedError

    else:
        raise ValueError

    pitch_mel = pitch_mel.float()

    if normalize_mean is not None:
        assert normalize_std is not None
        pitch_mel = normalize_pitch(pitch_mel, normalize_mean, normalize_std)

    return pitch_mel

def normalize_pitch(pitch, mean, std):
    zeros = (pitch == 0.0)
    pitch -= mean[:, None]
    pitch /= std[:, None]
    pitch[zeros] = 0.0
    return pitch

def normalize(S, max_abs_value=10., min_level_db = -12.):
    return torch.clamp((2 * max_abs_value) * ((S - min_level_db) / (-min_level_db)) - max_abs_value, -max_abs_value, max_abs_value)

class TTSDataset(torch.utils.data.Dataset):
    """
        1) loads audio,text pairs
        2) normalizes text and converts them to sequences of one-hot vectors
        3) computes mel-spectrograms from audio files.
    """
    def __init__(self,
                 dataset_path,
                 audiopaths_and_text,
                 text_cleaners,
                 n_mel_channels,
                 symbol_set='olli',
                 p_arpabet=1.0,
                 n_speakers=1,
                 load_mel_from_disk=True,
                 load_pitch_from_disk=True,
                 pitch_mean= 174.87988567095618,          #214.72203,  # LJSpeech defaults
                 pitch_std= 105.8709371279809,           #65.72038,
                 max_wav_value=None,
                 sampling_rate=None,
                 filter_length=None,
                 hop_length=None,
                 win_length=None,
                 mel_fmin=None,
                 mel_fmax=None,
                 prepend_space_to_text=False,
                 append_space_to_text=False,
                 pitch_online_dir=None,
                 betabinomial_online_dir=None,
                 use_betabinomial_interpolator=True,
                 pitch_online_method='pyin',
                 **ignored):

        # Expect a list of filenames
        if type(audiopaths_and_text) is str:
            audiopaths_and_text = [audiopaths_and_text]

        self.dataset_path = dataset_path
        self.audiopaths_and_text = load_filepaths_and_text(
            dataset_path, audiopaths_and_text,
            has_speakers=(n_speakers > 1))
        self.load_mel_from_disk = load_mel_from_disk
        if not load_mel_from_disk:
            self.max_wav_value = max_wav_value
            self.sampling_rate = sampling_rate
            self.stft = layers.TacotronSTFT(
                filter_length, hop_length, win_length,
                n_mel_channels, sampling_rate, mel_fmin, mel_fmax)
        self.load_pitch_from_disk = load_pitch_from_disk

        self.prepend_space_to_text = prepend_space_to_text
        self.append_space_to_text = append_space_to_text

        assert p_arpabet == 0.0 or p_arpabet == 1.0, (
            'Only 0.0 and 1.0 p_arpabet is currently supported. '
            'Variable probability breaks caching of betabinomial matrices.')

        self.tp = TextProcessing(symbol_set, text_cleaners, p_arpabet=p_arpabet)
        self.n_speakers = n_speakers
        self.pitch_tmp_dir = pitch_online_dir
        self.f0_method = pitch_online_method
        self.betabinomial_tmp_dir = betabinomial_online_dir
        self.use_betabinomial_interpolator = use_betabinomial_interpolator

        if use_betabinomial_interpolator:
            self.betabinomial_interpolator = BetaBinomialInterpolator()

        expected_columns = (2 + int(load_pitch_from_disk) + (n_speakers > 1))

        assert not (load_pitch_from_disk and self.pitch_tmp_dir is not None)

        if len(self.audiopaths_and_text[0]) < expected_columns:
            raise ValueError(f'Expected {expected_columns} columns in audiopaths file. '
                             'The format is <mel_or_wav>|[<pitch>|]<text>[|<speaker_id>]')

        if len(self.audiopaths_and_text[0]) > expected_columns:
            print('WARNING: Audiopaths file has more columns than expected')

        to_tensor = lambda x: torch.Tensor([x]) if type(x) is float else x
        #self.pitch_scaler = StandardScaler()
        self.pitch_mean = to_tensor(pitch_mean)
        self.pitch_std = to_tensor(pitch_std)

    def __getitem__(self, index):
        # Separate filename and text
        if self.n_speakers > 1:
            audiopath, *extra, text, speaker = self.audiopaths_and_text[index]
            speaker = int(speaker)
        else:
            audiopath, *extra, text = self.audiopaths_and_text[index]
            speaker = None

        mel = self.get_mel(audiopath).T
        text = self.get_text(text)
        pitch = self.get_pitch(index, mel.size(-1))
        pitch = pitch[:, :mel.size(-1)]# normalize the pitch len ONLY needed with external preprocess dataset
        energy = torch.norm(mel.float(), dim=0, p=2)
        attn_prior = self.get_prior(index, mel.shape[1], text.shape[0])

        assert pitch.size(-1) == mel.size(-1), print("{}{}".format(pitch.size(-1), mel.size(-1)))

        # No higher formants?
        if len(pitch.size()) == 1:
            pitch = pitch[None, :]

        return (text, mel, len(text), pitch, energy, speaker, attn_prior,
                audiopath)

    def __len__(self):
        return len(self.audiopaths_and_text)

    def get_mel(self, filename):
        if not self.load_mel_from_disk:
            #return None
            audio, sampling_rate = load_wav_to_torch(filename)
            if sampling_rate != self.stft.sampling_rate:
                raise ValueError("{} SR doesn't match target {} SR".format(
                    sampling_rate, self.stft.sampling_rate))
            audio_norm = audio / self.max_wav_value
            audio_norm = audio_norm.unsqueeze(0)
            audio_norm = torch.autograd.Variable(audio_norm,
                                                 requires_grad=False)
            melspec = self.stft.mel_spectrogram(audio_norm)
            melspec = torch.squeeze(melspec, 0).T
            #print(melspec)
        else:
            prefix = "/data2/tuong/OLLI-SPEECH-1.8/LJSpeech_style/mels"
            filename = filename.split('/')[-1].split('.')[0] + '.pt'
            filepath = '/'.join([prefix, filename])
            melspec = torch.load(filepath)
            melspec.requires_grad = False
            #melspec = normalize(melspec.T)
            melspec = melspec.T
            #print(melspec[10])
            # assert melspec.size(0) == self.stft.n_mel_channels, (
            #     'Mel dimension mismatch: given {}, expected {}'.format(
            #         melspec.size(0), self.stft.n_mel_channels))

        return melspec

    def get_text(self, text):
        #print(text.split())
        before = len(text.split())
        text = self.tp.encode_text(text)
        #print(text)
        after = len(text)
        #assert before == after, "{}, {}".format(before, after)
        #print("="*50)
        space = [self.tp.encode_text(",")[0]]

        if self.prepend_space_to_text:
            text = space + text

        if self.append_space_to_text:
            text = text + space

        return torch.LongTensor(text)

    def get_prior(self, index, mel_len, text_len):

        if self.use_betabinomial_interpolator:
            return torch.from_numpy(self.betabinomial_interpolator(mel_len,
                                                                   text_len))

        if self.betabinomial_tmp_dir is not None:
            audiopath, *_ = self.audiopaths_and_text[index]
            fname = Path(audiopath).relative_to(self.dataset_path)
            fname = fname.with_suffix('.pt')
            cached_fpath = Path(self.betabinomial_tmp_dir, fname)

            if cached_fpath.is_file():
                return torch.load(cached_fpath)

        attn_prior = beta_binomial_prior_distribution(text_len, mel_len)

        if self.betabinomial_tmp_dir is not None:
            cached_fpath.parent.mkdir(parents=True, exist_ok=True)
            torch.save(attn_prior, cached_fpath)

        return attn_prior

    def get_pitch(self, index, mel_len=None):
        audiopath, *fields = self.audiopaths_and_text[index]

        if self.n_speakers > 1:
            spk = int(fields[-1])
        else:
            spk = 0

        if self.load_pitch_from_disk:
            pitchpath = fields[0]
            pitch = torch.load(pitchpath)
            #print(pitch)
            if self.pitch_mean is not None:
                assert self.pitch_std is not None
                pitch = normalize_pitch(pitch, self.pitch_mean, self.pitch_std)
            #print(pitch)
            #print("="*77)
            return pitch

        if self.pitch_tmp_dir is not None:
            fname = Path(audiopath).relative_to(self.dataset_path)
            fname_method = fname.with_suffix('.pt')
            cached_fpath = Path(self.pitch_tmp_dir, fname_method)
            if cached_fpath.is_file():
                return torch.load(cached_fpath)

        # No luck so far - calculate
        wav = audiopath
        if not wav.endswith('.wav'):
            wav = re.sub('/mels/', '/wavs/', wav)
            wav = re.sub('.pt$', '.wav', wav)

        pitch_mel = estimate_pitch(wav, mel_len, self.f0_method,
                                   self.pitch_mean, self.pitch_std)

        if self.pitch_tmp_dir is not None and not cached_fpath.is_file():
            cached_fpath.parent.mkdir(parents=True, exist_ok=True)
            torch.save(pitch_mel, cached_fpath)

        return pitch_mel

class TTSCollate:
    """Zero-pads model inputs and targets based on number of frames per step"""

    def __call__(self, batch):
        """Collate training batch from normalized text and mel-spec"""
        # Right zero-pad all one-hot text sequences to max input length
        input_lengths, ids_sorted_decreasing = torch.sort(
            torch.LongTensor([len(x[0]) for x in batch]),
            dim=0, descending=True)
        max_input_len = input_lengths[0]

        text_padded = torch.LongTensor(len(batch), max_input_len)
        text_padded.zero_()
        for i in range(len(ids_sorted_decreasing)):
            text = batch[ids_sorted_decreasing[i]][0]
            text_padded[i, :text.size(0)] = text

        # Right zero-pad mel-spec
        num_mels = batch[0][1].size(0)
        max_target_len = max([x[1].size(1) for x in batch])

        # Include mel padded and gate padded
        mel_padded = torch.FloatTensor(len(batch), num_mels, max_target_len)
        mel_padded.zero_()
        output_lengths = torch.LongTensor(len(batch))
        for i in range(len(ids_sorted_decreasing)):
            mel = batch[ids_sorted_decreasing[i]][1]
            mel_padded[i, :, :mel.size(1)] = mel
            output_lengths[i] = mel.size(1)

        n_formants = batch[0][3].shape[0]
        pitch_padded = torch.zeros(mel_padded.size(0), n_formants,
                                   mel_padded.size(2), dtype=batch[0][3].dtype)
        energy_padded = torch.zeros_like(pitch_padded[:, 0, :])

        for i in range(len(ids_sorted_decreasing)):
            pitch = batch[ids_sorted_decreasing[i]][3]
            energy = batch[ids_sorted_decreasing[i]][4]
            pitch_padded[i, :, :pitch.shape[1]] = pitch
            energy_padded[i, :energy.shape[0]] = energy

        if batch[0][5] is not None:
            speaker = torch.zeros_like(input_lengths)
            for i in range(len(ids_sorted_decreasing)):
                speaker[i] = batch[ids_sorted_decreasing[i]][5]
        else:
            speaker = None

        attn_prior_padded = torch.zeros(len(batch), max_target_len,
                                        max_input_len)
        attn_prior_padded.zero_()
        for i in range(len(ids_sorted_decreasing)):
            prior = batch[ids_sorted_decreasing[i]][6]
            attn_prior_padded[i, :prior.size(0), :prior.size(1)] = prior

        # Count number of items - characters in text
        len_x = [x[2] for x in batch]
        len_x = torch.Tensor(len_x)

        audiopaths = [batch[i][7] for i in ids_sorted_decreasing]

        return (text_padded, input_lengths, mel_padded, output_lengths, len_x,
                pitch_padded, energy_padded, speaker, attn_prior_padded,
                audiopaths)

def batch_to_gpu(batch):
    (text_padded, input_lengths, mel_padded, output_lengths, len_x,
     pitch_padded, energy_padded, speaker, attn_prior, audiopaths) = batch

    text_padded = to_gpu(text_padded).long()
    input_lengths = to_gpu(input_lengths).long()
    mel_padded = to_gpu(mel_padded).float()
    output_lengths = to_gpu(output_lengths).long()
    pitch_padded = to_gpu(pitch_padded).float()
    energy_padded = to_gpu(energy_padded).float()
    attn_prior = to_gpu(attn_prior).float()
    if speaker is not None:
        speaker = to_gpu(speaker).long()

    # Alignments act as both inputs and targets - pass shallow copies
    x = [text_padded, input_lengths, mel_padded, output_lengths,
         pitch_padded, energy_padded, speaker, attn_prior, audiopaths]
    y = [mel_padded, input_lengths, output_lengths]
    len_x = torch.sum(output_lengths)
    return (x, y, len_x)


import shutil
import os

# This is not compatible with astropy.wcs and astropy.pyfits
try:

    import pyfits
    import pywcs

except ImportError:

    # A dummy version that doesn't do anything. The cutout is not supported
    # if pyfits and pywcs are not installed
    def cutout(filename, ra_or_l, dec_or_b, coordsys, radius, outfile, clobber=True):

        if os.path.exists(outfile):

            os.remove(outfile)

        shutil.copy(filename, outfile)

else:

    import numpy
    from pyLikelihood import SkyDir
    from GtBurst.angularDistance import getAngularDistance
    from GtBurst.angularDistance import getBoundingCoordinates

    class DimensionError(ValueError):
        pass

    def cutout(filename, ra_or_l, dec_or_b, coordsys, radius, outfile, clobber=True):
        """
        Inputs:
            file  - .fits filename or pyfits HDUList. If HDU is 3d (data-cube), the 3rd dimension
                    (the one which is not a sky coordinate) will be kept untouched. Dimensions > 3
                    are not supported
            ra_or_l,dec_or_b - Longitude and latitude of the center of the new image. The longitude
                               coordinate (R.A. or L) goes from 0 to 360, while the latitude coordinate
                               goes from -90 to 90.
            coordsys - Coordinate system for the center of the new image ('galactic' or 'equatorial')
            radius - radius of the region of interest (deg)
            outfile - output file
            clobber - overwrite the file if existing? (True or False)
        """

        if coordsys not in ['equatorial','galactic']:
          raise ValueError("Unknown coordinate system '%s'" %(coordsys))

        if(ra_or_l < 0 or ra_or_l > 360):
          raise RuntimeError("The longitude coordinate must be 0 < longitude < 360")

        if(dec_or_b < -90 or dec_or_b > 90):
          raise RuntimeError("The latitude coordinate must be -90 < latitude < 90")

        with pyfits.open(filename) as f:

          if(f[0].data.ndim==3):
            isCube                     = True
          elif(f[0].data.ndim==2):
            isCube                     = False
          else:
            raise RuntimeError("Do not know how to handle a cube with %s dimensions" %(f[0].data.shape[0]))
          pass

          head                         = f[0].header.copy()

          cd1                          = head.get('CDELT1') if head.get('CDELT1') else head.get('CD1_1')
          cd2                          = head.get('CDELT2') if head.get('CDELT2') else head.get('CD2_2')
          if cd1 is None or cd2 is None:
              raise Exception("Missing CD or CDELT keywords in header")

          wcs                         = pywcs.WCS(head)

          #Ensure that the center is expressed in the same coordinate system as the original
          #image
          if coordsys=='equatorial' and wcs.wcs.lngtyp=='GLON':

                #Convert RA,Dec in Galactic coordinates
                sdir                        = SkyDir(ra_or_l,dec_or_b,SkyDir.EQUATORIAL)
                xc,yc                       = (sdir.l(),sdir.b())

          elif coordsys=='galactic' and wcs.wcs.lngtyp=='RA':

                #Convert L,B in Equatorial coordinates
                sdir                        = SkyDir(ra_or_l,dec_or_b,SkyDir.EQUATORIAL)
                xc,yc                       = (sdir.ra(),sdir.dec())

          else:

               #Image and input are in the same system already
               xc,yc                        = (ra_or_l,dec_or_b)

          pass

          #Find the pixel corresponding to the center
          if(isCube):
            #Data cube
            coord                      = numpy.array([[xc],[yc],[1]]).T
            xx,yy,z                    = wcs.wcs_sky2pix(coord,0)[0]
            shapez,shapey,shapex       = f[0].data.shape

            #Compute the sky coordinates of all pixels
            #(will use them later for the angular distance)

            #The code below is much faster, but does the same thing as this
            #one here:

            #coord                        = numpy.zeros((shapex*shapey,3))
            #h                            = 0
            #for i in range(shapex):
            #  for j in range(shapey):
            #    coord[h]                 = [i+1,j+1,1]
            #    h                       += 1

            coord                        = numpy.ones((shapex*shapey,3))
            firstColBase                 = numpy.arange(shapex)+1
            firstColumn                  = numpy.repeat(firstColBase,shapey)
            secondColumn                 = numpy.array(list(range(shapey))*shapex)+1
            coord[:,0]                   = firstColumn
            coord[:,1]                   = secondColumn

            #Note that pix2sky always return the latitude from 0 to 360 deg
            res                          = wcs.wcs_pix2sky(coord,1)
            ras                          = res[:,0]
            decs                         = res[:,1]

          else:
            #Normal image
            coord                       = numpy.array([[xc],[yc]]).T
            xx,yy                       = wcs.wcs_sky2pix(coord,0)[0]
            shapey,shapex               = f[0].data.shape

            #Compute the sky coordinates of all pixels
            #(will use them later for the angular distance)
            coord                        = numpy.ones((shapex*shapey,2))
            firstColBase                 = numpy.arange(shapex)+1
            firstColumn                  = numpy.repeat(firstColBase,shapey)
            secondColumn                 = numpy.array(list(range(shapey))*shapex)+1
            coord[:,0]                   = firstColumn
            coord[:,1]                   = secondColumn

            #Note that pix2sky always return the latitude from 0 to 360 deg
            res                          = wcs.wcs_pix2sky(coord,1)
            ras                          = res[:,0]
            decs                         = res[:,1]
          pass

          print(("Center is (%s,%s) pixel, (%s,%s) sky" %(xx,yy,xc,yc)))

          #Cannot deal with fractional pixel
          if(xx - int(xx) >= 1e-4):
            print(("Approximating the X pixel: %s -> %s" %(xx,int(xx))))
            xx                          = int(xx)
          if(yy - int(yy) >= 1e-4):
            print(("Approximating the Y pixel: %s -> %s" %(yy,int(yy))))
            yy                          = int(yy)
          pass

          #Now select the pixels to keep
          #Pre-select according to a bounding box
          #(huge gain of speed in the computation of distances)
          ra_min_,ra_max_,dec_min_,dec_max_,pole = getBoundingCoordinates(xc,yc,radius)

          if(pole):
            #Nothing we can really do, except masking out (zeroing) all the useless parts
            print("\nOne of the poles is within the region. Cannot cut the image. I will zero-out useless parts")
            img                           = f[0].data

            #Find the pixel corresponding to (xc,yc-radius)
            coord                         = numpy.array([[xc],[yc-radius],[1]]).T
            res                           = wcs.wcs_sky2pix(coord,0)[0]
            mask_ymax                     = int(numpy.ceil(res[1]))

            #Now find the pixel corresponding to (xc,yc+radius)
            coord                         = numpy.array([[xc],[yc+radius-180.0],[1]]).T
            res                           = wcs.wcs_sky2pix(coord,0)[0]
            mask_ymin                     = int(numpy.floor(res[1]))

            img[:,mask_ymin:mask_ymax,:]  = img[:,mask_ymin:mask_ymax,:]*0.0

          else:
            if(ra_min_ > ra_max_):
              #Circular inversion (ex: ra_min_ = 340.0 and ra_max_ = 20.0)
              idx                         = (ra_min_ <= ras) | (ras <= ra_max_)
            else:
              idx                         = (ra_min_ <= ras) & (ras <= ra_max_)
            pass

            if(dec_min_ > dec_max_):
              #Circular inversion (ex: ra_min_ = 340.0 and ra_max_ = 20.0)
              idx                         = ((dec_min_ <= decs) | (decs <= dec_max_)) & idx
            else:
              idx                         = ((dec_min_ <= decs) & (decs <= dec_max_)) & idx
            pass

            ras                           = ras[idx]
            decs                          = decs[idx]

            #Compute all angular distances of remaining pixels
            distances                     = getAngularDistance(xc,yc,ras,decs)

            #Select all pixels within the provided radius
            idx                           = (distances <= radius)
            selected_ras                  = ras[idx]
            selected_decs                 = decs[idx]

            #Now transform back into pixels values
            if(isCube):

              coord                      = numpy.vstack([selected_ras,selected_decs,[1]*selected_ras.shape[0]]).T

            else:

              coord                      = numpy.vstack([selected_ras,selected_decs]).T

            pass

            res                        = wcs.wcs_sky2pix(coord,0)

            #Now check if the range of x is not continuous (i.e., we are
            #wrapping around the borders)
            uniquex                    = numpy.unique(res[:,0])
            deltas                     = uniquex[1:]-uniquex[:-1]
            if(deltas.max()>1):
              #We are wrapping around the borders
              # |-------x1             x2-------|
              #We want to express x2 as a negative index starting
              #from the right border, and set it as xmin.
              #Then we set x1 as xmax.
              #This way the .take method below will start accumulating
              #the image from x2 to the right border |, then from the left
              #border | to x1, in this order
              #Find x2
              x2id                     = deltas.argmax()+1
              x2                       = int(uniquex[x2id])
              x1                       = int(uniquex[x2id-1])
              xmin                     = x2-shapex
              xmax                     = x1

              ymin,ymax                = (int(res[:,1].min()),int(res[:,1].max()))

            else:

              xmin,xmax,ymin,ymax      = (int(res[:,0].min()),int(res[:,0].max()),int(res[:,1].min()),int(res[:,1].max()))

            pass

            print(("X range -> %s - %s" %(xmin,xmax)))
            print(("Y range -> %s - %s" %(ymin,ymax)))
            print(("Input image shape is ([z],y,x) = %s" %(str(f[0].shape))))

            #Using the mode='wrap' option we wrap around the edges of the image,
            #if ymin is negative
            if(isCube):

              img                        = f[0].data.take(list(range(ymin,ymax)),mode='wrap', axis=1).take(list(range(xmin,xmax)),mode='wrap',axis=2)

            else:

              img                        = f[0].data.take(list(range(ymin,ymax)),mode='wrap', axis=0).take(list(range(xmin,xmax)),mode='wrap',axis=1)

            pass

            #Put the origin of the projection in the right place
            #in the new image
            head['CRPIX1']               -= xmin
            head['CRPIX2']               -= ymin

            #Update the length of the axis
            head['NAXIS1']               = int(xmax-xmin)
            head['NAXIS2']               = int(ymax-ymin)

            if head.get('NAXIS1') == 0 or head.get('NAXIS2') == 0:
                raise ValueError("Map has a 0 dimension: %i,%i." % (head.get('NAXIS1'),head.get('NAXIS2')))
          pass

          newfile = pyfits.PrimaryHDU(data=img,header=head)

          newfile.writeto(outfile,clobber=clobber)
          #Append the other extension, if present
          for i in range(1,len(f)):
            pyfits.append(outfile,f[i].data,header=f[i].header)

        pass #Close the input file

        #Now re-open the output file and fix the wcs
        #by moving the reference pixel to the 1,1 pixel
        #This guarantee that no pixel will be at a distance larger than 180 deg
        #from the reference pixel, which would confuse downstream software
        if(not pole):
          with pyfits.open(outfile,'update') as outf:
            head                         = outf[0].header
            #Get the
            newwcs                          = pywcs.WCS(head)

            #Find the sky coordinates of the 1,1 pixel
            if(isCube):
              #Data cube
              coord                       = numpy.array([[1],[1],[1]]).T
              sx,sy,z                     = newwcs.wcs_pix2sky(coord,1)[0]
            else:
              #Normal image
              coord                       = numpy.array([[1],[1]]).T
              sx,sy                       = newwcs.wcs_pix2sky(coord,1)[0]
            pass

            head['CRPIX1']                = 1
            head['CRVAL1']                = sx

        pass


from __future__ import print_function, absolute_import, division

import contextlib
import warnings
from copy import deepcopy

try:
    import builtins
except ImportError:
    # python2
    import __builtin__ as builtins

import dask.array as da
import numpy as np
from astropy.wcs.utils import proj_plane_pixel_area
from astropy.wcs import (WCSSUB_SPECTRAL, WCSSUB_LONGITUDE, WCSSUB_LATITUDE)
from . import wcs_utils
from .utils import FITSWarning, AstropyUserWarning, WCSCelestialError
from astropy import log
from astropy.io import fits
from astropy.wcs.utils import is_proj_plane_distorted
from astropy.io.fits import BinTableHDU, Column
from astropy import units as u
import itertools
import re
from radio_beam import Beam

def _fix_spectral(wcs):
    """
    Attempt to fix a cube with an invalid spectral axis definition.  Only uses
    well-known exceptions, e.g. CTYPE = 'VELOCITY'.  For the rest, it will try
    to raise a helpful error.
    """

    axtypes = wcs.get_axis_types()

    types = [a['coordinate_type'] for a in axtypes]

    if wcs.naxis not in (3, 4):
        raise TypeError("The WCS has {0} axes of types {1}".format(len(types),
                                                                   types))

    # sanitize noncompliant headers
    if 'spectral' not in types:
        log.warning("No spectral axis found; header may be non-compliant.")
        for ind,tp in enumerate(types):
            if tp not in ('celestial','stokes'):
                if wcs.wcs.ctype[ind] in wcs_utils.bad_spectypes_mapping:
                    wcs.wcs.ctype[ind] = wcs_utils.bad_spectypes_mapping[wcs.wcs.ctype[ind]]

    return wcs

def _split_stokes(array, wcs):
    """
    Given a 4-d data cube with 4-d WCS (spectral cube + stokes) return a
    dictionary of data and WCS objects for each Stokes component

    Parameters
    ----------
    array : `~numpy.ndarray`
        The input 3-d array with two position dimensions, one spectral
        dimension, and a Stokes dimension.
    wcs : `~astropy.wcs.WCS`
        The input 3-d WCS with two position dimensions, one spectral
        dimension, and a Stokes dimension.
    """

    if array.ndim not in (3,4):
        raise ValueError("Input array must be 3- or 4-dimensional for a"
                         " STOKES cube")

    if wcs.wcs.naxis != 4:
        raise ValueError("Input WCS must be 4-dimensional for a STOKES cube")

    wcs = _fix_spectral(wcs)

    # reverse from wcs -> numpy convention
    axtypes = wcs.get_axis_types()[::-1]

    types = [a['coordinate_type'] for a in axtypes]

    try:
        # Find stokes dimension
        stokes_index = types.index('stokes')
    except ValueError:
        # stokes not in list, but we are 4d
        if types.count('celestial') == 2 and types.count('spectral') == 1:
            if None in types:
                stokes_index = types.index(None)
                log.warning("FITS file has no STOKES axis, but it has a blank"
                            " axis type at index {0} that is assumed to be "
                            "stokes.".format(4-stokes_index))
            else:
                for ii,tp in enumerate(types):
                    if tp not in ('celestial', 'spectral'):
                        stokes_index = ii
                        stokes_type = tp

                log.warning("FITS file has no STOKES axis, but it has an axis"
                            " of type {1} at index {0} that is assumed to be "
                            "stokes.".format(4-stokes_index, stokes_type))
        else:
            raise IOError("There are 4 axes in the data cube but no STOKES "
                          "axis could be identified")

    # TODO: make the stokes names more general
    stokes_names = ["I", "Q", "U", "V"]

    stokes_arrays = {}

    wcs_slice = wcs_utils.drop_axis(wcs, wcs.naxis - 1 - stokes_index)

    if array.ndim == 4:
        for i_stokes in range(array.shape[stokes_index]):

            array_slice = [i_stokes if idim == stokes_index else slice(None)
                           for idim in range(array.ndim)]

            stokes_arrays[stokes_names[i_stokes]] = array[tuple(array_slice)]
    else:
        # 3D array with STOKES as a 4th header parameter
        stokes_arrays['I'] = array

    return stokes_arrays, wcs_slice

def _orient(array, wcs):
    """
    Given a 3-d spectral cube and WCS, swap around the axes so that the
    spectral axis cube is the first in Numpy notation, and the last in WCS
    notation.

    Parameters
    ----------
    array : `~numpy.ndarray`
        The input 3-d array with two position dimensions and one spectral
        dimension.
    wcs : `~astropy.wcs.WCS`
        The input 3-d WCS with two position dimensions and one spectral
        dimension.
    """

    if array.ndim != 3:
        raise ValueError("Input array must be 3-dimensional")

    if wcs.wcs.naxis != 3:
        raise ValueError("Input WCS must be 3-dimensional")

    wcs = wcs_utils.diagonal_wcs_to_cdelt(_fix_spectral(wcs))

    # reverse from wcs -> numpy convention
    axtypes = wcs.get_axis_types()[::-1]

    types = [a['coordinate_type'] for a in axtypes]

    n_celestial = types.count('celestial')

    if n_celestial == 0:
        raise ValueError('No celestial axes found in WCS')
    elif n_celestial != 2:
        raise ValueError('WCS should contain 2 celestial dimensions but '
                         'contains {0}'.format(n_celestial))

    n_spectral = types.count('spectral')

    if n_spectral == 0:
        raise ValueError('No spectral axes found in WCS')
    elif n_spectral != 1:
        raise ValueError('WCS should contain one spectral dimension but '
                         'contains {0}'.format(n_spectral))

    nums = [None if a['coordinate_type'] != 'celestial' else a['number']
            for a in axtypes]

    if 'stokes' in types:
        raise ValueError("Input WCS should not contain stokes")

    t = [types.index('spectral'), nums.index(1), nums.index(0)]
    if t == [0, 1, 2]:
        result_array = array
    else:
        result_array = array.transpose(t)

    result_wcs = wcs.sub([WCSSUB_LONGITUDE, WCSSUB_LATITUDE, WCSSUB_SPECTRAL])

    return result_array, result_wcs

def slice_syntax(f):
    """
    This decorator wraps a function that accepts a tuple of slices.

    After wrapping, the function acts like a property that accepts
    bracket syntax (e.g., p[1:3, :, :])

    Parameters
    ----------
    f : function
    """

    def wrapper(self):
        result = SliceIndexer(f, self)
        result.__doc__ = f.__doc__
        return result

    wrapper.__doc__ = slice_doc.format(f.__doc__ or '',
                                       f.__name__)

    result = property(wrapper)

    return result

slice_doc = """
{0}

Notes
-----
Supports efficient Numpy slice notation,
like ``{1}[0:3, :, 2:4]``
"""

class SliceIndexer(object):

    def __init__(self, func, _other):
        self._func = func
        self._other = _other

    def __getitem__(self, view):
        result = self._func(self._other, view)
        if isinstance(result, da.Array):
            result = result.compute()
        return result

    @property
    def size(self):
        return self._other.size

    @property
    def ndim(self):
        return self._other.ndim

    @property
    def shape(self):
        return self._other.shape

    def __iter__(self):
        raise Exception("You need to specify a slice (e.g. ``[:]`` or "
                        "``[0,:,:]`` in order to access this property.")

# TODO: make this into a proper configuration item
# TODO: make threshold depend on memory?
MEMORY_THRESHOLD=1e8

def is_huge(cube):
    if cube.size < MEMORY_THRESHOLD:  # smallish
        return False
    else:
        return True

def iterator_strategy(cube, axis=None):
    """
    Guess the most efficient iteration strategy
    for iterating over a cube, given its size and layout

    Parameters
    ----------
    cube : SpectralCube instance
        The cube to iterate over
    axis : [0, 1, 2]
        For reduction methods, the axis that is
        being collapsed

    Returns
    -------
    strategy : ['cube' | 'ray' | 'slice']
        The recommended iteration strategy.
        *cube* recommends working with the entire array in memory
        *slice* recommends working with one slice at a time
        *ray*  recommends working with one ray at a time
    """
    # pretty simple for now
    if cube.size < 1e8:  # smallish
        return 'cube'
    return 'slice'

def try_load_beam(header):
    '''
    Try loading a beam from a FITS header.
    '''

    try:
        beam = Beam.from_fits_header(header)
        return beam
    except Exception as ex:
        # We don't emit a warning if no beam was found since it's ok for
        # cubes to not have beams
        # if 'No BMAJ' not in str(ex):
        #     warnings.warn("Could not parse beam information from header."
        #                   "  Exception was: {0}".format(ex.__repr__()),
        #                   FITSWarning
        #                  )

        # Avoid warning since cubes don't have a beam
        # Warning now provided when `SpectralCube.beam` is None
        beam = None

    return beam

def try_load_beams(data):
    '''
    Try loading a beam table from a FITS HDU list.
    '''
    try:
        from radio_beam import Beam
    except ImportError:
        warnings.warn("radio_beam is not installed. No beam "
                      "can be created.",
                      ImportError
                     )

    if isinstance(data, fits.BinTableHDU):
        if 'BPA' in data.data.names:
            beam_table = data.data
            return beam_table
        else:
            raise ValueError("No beam table found")
    elif isinstance(data, fits.HDUList):

        for ihdu, hdu_item in enumerate(data):
            if isinstance(hdu_item, (fits.PrimaryHDU, fits.ImageHDU)):
                beam = try_load_beams(hdu_item.header)
            elif isinstance(hdu_item, fits.BinTableHDU):
                if 'BPA' in hdu_item.data.names:
                    beam_table = hdu_item.data
                    return beam_table

        try:
            # if there was a beam in a header, but not a beam table
            return beam
        except NameError:
            # if the for loop has completed, we didn't find a beam table
            raise ValueError("No beam table found")
    elif isinstance(data, (fits.PrimaryHDU, fits.ImageHDU)):
        return try_load_beams(data.header)
    elif isinstance(data, fits.Header):
        try:
            beam = Beam.from_fits_header(data)
            return beam
        except Exception as ex:
            # warnings.warn("Could not parse beam information from header."
            #               "  Exception was: {0}".format(ex.__repr__()),
            #               FITSWarning
            #              )

            # Avoid warning since cubes don't have a beam
            # Warning now provided when `SpectralCube.beam` is None
            beam = None
    else:
        raise ValueError("How did you get here?  This is some sort of error.")

def beams_to_bintable(beams):
    """
    Convert a list of beams to a CASA-style BinTableHDU
    """

    c1 = Column(name='BMAJ', format='1E', array=[bm.major.to(u.arcsec).value for bm in beams], unit=u.arcsec.to_string('FITS'))
    c2 = Column(name='BMIN', format='1E', array=[bm.minor.to(u.arcsec).value for bm in beams], unit=u.arcsec.to_string('FITS'))
    c3 = Column(name='BPA', format='1E', array=[bm.pa.to(u.deg).value for bm in beams], unit=u.deg.to_string('FITS'))
    #c4 = Column(name='CHAN', format='1J', array=[bm.meta['CHAN'] if 'CHAN' in bm.meta else 0 for bm in beams])
    c4 = Column(name='CHAN', format='1J', array=np.arange(len(beams)))
    c5 = Column(name='POL', format='1J', array=[bm.meta['POL'] if 'POL' in bm.meta else 0 for bm in beams])

    bmhdu = BinTableHDU.from_columns([c1, c2, c3, c4, c5])
    bmhdu.header['EXTNAME'] = 'BEAMS'
    bmhdu.header['EXTVER'] = 1
    bmhdu.header['XTENSION'] = 'BINTABLE'
    bmhdu.header['NCHAN'] = len(beams)
    bmhdu.header['NPOL'] = len(set([bm.meta['POL'] for bm in beams if 'POL' in bm.meta]))
    return bmhdu

def beam_props(beams, includemask=None):
    '''
    Returns separate quantities for the major, minor, and PA of a list of
    beams.
    '''
    if includemask is None:
        includemask = itertools.cycle([True])

    major = u.Quantity([bm.major for bm, incl in zip(beams, includemask)
                        if incl], u.deg)
    minor = u.Quantity([bm.minor for bm, incl in zip(beams, includemask)
                        if incl], u.deg)
    pa = u.Quantity([bm.pa for bm, incl in zip(beams, includemask)
                     if incl], u.deg)

    return major, minor, pa

def largest_beam(beams, includemask=None):
    """
    Returns the largest beam (by area) in a list of beams.
    """

    from radio_beam import Beam

    major, minor, pa = beam_props(beams, includemask)
    largest_idx = (major * minor).argmax()
    new_beam = Beam(major=major[largest_idx], minor=minor[largest_idx],
                    pa=pa[largest_idx])

    return new_beam

def smallest_beam(beams, includemask=None):
    """
    Returns the smallest beam (by area) in a list of beams.
    """

    from radio_beam import Beam

    major, minor, pa = beam_props(beams, includemask)
    smallest_idx = (major * minor).argmin()
    new_beam = Beam(major=major[smallest_idx], minor=minor[smallest_idx],
                    pa=pa[smallest_idx])

    return new_beam

@contextlib.contextmanager
def _map_context(numcores):
    """
    Mapping context manager to allow parallel mapping or regular mapping
    depending on the number of cores specified.

    The builtin map is overloaded to handle python3 problems: python3 returns a
    generator, while ``multiprocessing.Pool.map`` actually runs the whole thing
    """
    if numcores is not None and numcores > 1:
        try:
            from joblib import Parallel, delayed
            from joblib.pool import has_shareable_memory
            map = lambda x,y: Parallel(n_jobs=numcores)(delayed(has_shareable_memory)(x))(y)
            parallel = True
        except ImportError:
            map = lambda x,y: list(builtins.map(x,y))
            warnings.warn("Could not import joblib.  "
                          "map will be non-parallel.",
                          ImportError
                         )
            parallel = False
    else:
        parallel = False
        map = lambda x,y: list(builtins.map(x,y))

    yield map

def convert_bunit(bunit):
    '''
    Convert a BUNIT string to a quantity

    Parameters
    ----------
    bunit : str
        String to convert to an `~astropy.units.Unit`

    Returns
    -------
    unit : `~astropy.unit.Unit`
        Corresponding unit.
    '''

    # special case: CASA (sometimes) makes non-FITS-compliant jy/beam headers
    bunit_lower = re.sub(r"\s", "", bunit.lower())
    if bunit_lower == 'jy/beam':
        unit = u.Jy / u.beam
    else:
        try:
            unit = u.Unit(bunit)
        except ValueError:
            warnings.warn("Could not parse unit {0}".format(bunit),
                          AstropyUserWarning)
            unit = None

    return unit

def world_take_along_axis(cube, position_plane, axis):
    '''
    Convert a 2D plane of pixel positions to the equivalent WCS coordinates.
    For example, this will convert `argmax`
    along the spectral axis to the equivalent spectral value (e.g., velocity at
    peak intensity).

    Parameters
    ----------
    cube : SpectralCube
        A spectral cube.
    position_plane : 2D numpy.ndarray
        2D array of pixel positions along `axis`. For example, `position_plane` can
        be the output of `argmax` or `argmin` along an axis.
    axis : int
        The axis that `position_plane` is collapsed along.

    Returns
    -------
    out : astropy.units.Quantity
        2D array of WCS coordinates.
    '''

    if wcs_utils.is_pixel_axis_to_wcs_correlated(cube.wcs, axis):
        raise WCSCelestialError("world_take_along_axis requires the celestial axes"
                                " to be aligned along image axes.")

    # Get 1D slice along that axis.
    world_slice = [0, 0]
    world_slice.insert(axis, slice(None))

    world_coords = cube.world[tuple(world_slice)][axis]

    world_newaxis = [np.newaxis] * 2
    world_newaxis.insert(axis, slice(None))
    world_newaxis = tuple(world_newaxis)

    plane_newaxis = [slice(None), slice(None)]
    plane_newaxis.insert(axis, np.newaxis)
    plane_newaxis = tuple(plane_newaxis)

    out = np.take_along_axis(world_coords[world_newaxis],
                             position_plane[plane_newaxis], axis=axis)
    out = out.squeeze()

    return out

def _has_beam(obj):
    if hasattr(obj, '_beam'):
        return obj._beam is not None
    else:
        return False

def _has_beams(obj):
    if hasattr(obj, '_beams'):
        return obj._beams is not None
    else:
        return False

def bunit_converters(obj, unit, equivalencies=(), freq=None):
    '''
    Handler for all brightness unit conversions, including: K, Jy/beam, Jy/pix, Jy/sr.
    This also includes varying resolution spectral cubes, where the beam size varies along
    the frequency axis.

    Parameters
    ----------
    obj : {SpectralCube, LowerDimensionalObject}
        A spectral cube or any other lower dimensional object.
    unit : `~astropy.units.Unit`
        Unit to convert `obj` to.
    equivalencies : tuple, optional
        Initial list of equivalencies.
    freq : `~astropy.unit.Quantity`, optional
        Frequency to use for spectral conversions. If the spectral axis is available, the
        frequencies will already be defined.

    Outputs
    -------
    factor : `~numpy.ndarray`
        Array of factors for the unit conversion.

    '''

    # Add a simple check it the new unit is already equivalent, and so we don't need
    # any additional unit equivalencies
    if obj.unit.is_equivalent(unit):
        # return equivalencies
        factor = obj.unit.to(unit, equivalencies=equivalencies)
        return np.array([factor])

    # Determine the bunit "type". This will determine what information we need for the unit conversion.
    has_btemp = obj.unit.is_equivalent(u.K) or unit.is_equivalent(u.K)
    has_perbeam = obj.unit.is_equivalent(u.Jy/u.beam) or unit.is_equivalent(u.Jy/u.beam)
    has_perangarea = obj.unit.is_equivalent(u.Jy/u.sr) or unit.is_equivalent(u.Jy/u.sr)
    has_perpix = obj.unit.is_equivalent(u.Jy/u.pix) or unit.is_equivalent(u.Jy/u.pix)

    # Is there any beam object defined?
    has_beam = _has_beam(obj) or _has_beams(obj)

    # Set if this is a varying resolution object
    has_beams = _has_beams(obj)

    # Define freq, if needed:
    if any([has_perangarea, has_perbeam, has_btemp]):
        # Create a beam equivalency for brightness temperature
        # This requires knowing the frequency along the spectral axis.
        if freq is None:
            try:
                freq = obj.with_spectral_unit(u.Hz).spectral_axis
            except AttributeError:
                raise TypeError("Object of type {0} has no spectral "
                                "information. `freq` must be provided for"
                                " unit conversion from Jy/beam"
                                .format(type(obj)))
        else:
            if not freq.unit.is_equivalent(u.Hz):
                raise u.UnitsError("freq must be given in equivalent "
                                   "frequency units.")

            freq = freq.reshape((-1,))

    else:
        freq = [None]

    # To handle varying resolution objects, loop through "channels"
    # Default to a single iteration for a 2D spatial object or when a beam is not defined
    # This allows handling all 1D, 2D, and 3D data products.
    if has_beams:
        iter = range(len(obj.beams))
        beams = obj.beams
    elif has_beam:
        iter = range(0, 1)
        beams = [obj.beam]
    else:
        iter = range(0, 1)
        beams = [None]

    # Append the unit conversion factors
    factors = []

    # Iterate through spectral channels.
    for ii in iter:

        beam = beams[ii]

        # Use the range of frequencies when the beam does not change. Otherwise, select the
        # frequency corresponding to this beam.
        if has_beams:
            thisfreq = freq[ii]
        else:
            thisfreq = freq

        # Changes in beam require a new equivalency for each.
        this_equivalencies = deepcopy(equivalencies)

        # Equivalencies for Jy per ang area.
        if has_perangarea:
            bmequiv_angarea = u.brightness_temperature(thisfreq)

            this_equivalencies = list(this_equivalencies) + bmequiv_angarea

        # Beam area equivalencies for Jy per beam and/or Jy per ang area
        if has_perbeam or has_perangarea:
            if not has_beam:
                raise ValueError("To convert cubes with Jy/beam units, "
                                "the cube needs to have a beam defined.")

            # create a beam equivalency for brightness temperature
            bmequiv = beam.jtok_equiv(thisfreq)

            # NOTE: `beamarea_equiv` was included in the radio-beam v0.3.3 release
            # The if/else here handles potential cases where earlier releases are installed.
            if hasattr(beam, 'beamarea_equiv'):
                bmarea_equiv = beam.beamarea_equiv
            else:
                bmarea_equiv = u.beam_angular_area(beam.sr)

            this_equivalencies = list(this_equivalencies) + bmequiv + bmarea_equiv

        # Equivalencies for Jy per pixel area.
        if has_perpix:

            if not obj.wcs.has_celestial:
                raise ValueError("Spatial WCS information is required for unit conversions"
                                " involving spatial areas (e.g., Jy/pix, Jy/sr)")

            pix_area = (proj_plane_pixel_area(obj.wcs.celestial) * u.deg**2).to(u.sr)

            pix_area_equiv = [(u.Jy / u.pix, u.Jy / u.sr,
                            lambda x: x / pix_area.value,
                            lambda x: x * pix_area.value)]

            this_equivalencies = list(this_equivalencies) + pix_area_equiv

            # Define full from brightness temp to Jy / pix.
            # Otherwise isn't working in 1 step
            if has_btemp:
                if not has_beam:
                    raise ValueError("Conversions between K and Jy/beam or Jy/pix"
                                    "requires the cube to have a beam defined.")

                jtok_factor = beam.jtok(thisfreq) / (u.Jy / u.beam)

                # We're going to do this piecemeal because it's easier to conceptualize
                # We specifically anchor these conversions based on the beam area. So from
                # beam to pix, this is beam -> angular area -> area per pixel
                # Altogether:
                # K ->  Jy/beam -> Jy /sr - > Jy / pix
                forward_factor = 1 / (jtok_factor * (beam.sr / u.beam) / (pix_area / u.pix))
                reverse_factor = jtok_factor * (beam.sr / u.beam) / (pix_area / u.pix)

                pix_area_btemp_equiv = [(u.K, u.Jy / u.pix,
                                        lambda x: x * forward_factor.value,
                                        lambda x: x * reverse_factor.value)]

                this_equivalencies = list(this_equivalencies) + pix_area_btemp_equiv

            # Equivalencies between pixel and angular areas.
            if has_perbeam:
                if not has_beam:
                    raise ValueError("Conversions between Jy/beam or Jy/pix"
                                    "requires the cube to have a beam defined.")

                beam_area = beam.sr

                pix_area_btemp_equiv = [(u.Jy / u.pix, u.Jy / u.beam,
                                        lambda x: x * (beam_area / pix_area).value,
                                        lambda x: x * (pix_area / beam_area).value)]

                this_equivalencies = list(this_equivalencies) + pix_area_btemp_equiv

        factor = obj.unit.to(unit, equivalencies=this_equivalencies)
        factors.append(factor)

    if has_beams:
        return factors
    else:
        # Slice along first axis to return a 1D array.
        return factors[0]


#  Import necessary packages
import logging
import random
from time import sleep

import numpy as np

from pymeasure.experiment import IntegerParameter, FloatParameter, Parameter
from pymeasure.experiment import Procedure
from pymeasure.instruments.keithley import Keithley2600

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())

##########
# global #
##########

##########################
# Pre-Defined Procedures #
##########################

# todo: Define Standard Procedures

###########################
# User-Defined Procedures #
###########################
# Todo: Write Sweep Procedure
# Todo: Write Cycle Procedure
# Todo: Write Pulse Procedure (Just factory script)
# Todo: Write Pulse Procedure (write Pulse, read cycle)
# Todo: Write Retention Procedure
# Todo: Write Endurance Procedure

class ProcedureWithInstruments(Procedure):
    """Class that adds a dict of instruments to Procedure. SelectionWindow uses this dict to pass
    user selected, connected devices (names & adresses)"""
    def __init__(self, instruments_dict={}):
        Procedure.__init__(self)
        instrument_adress = ''

class IVCycles(ProcedureWithInstruments):

    #define required instruments
    required_instruments = ['Keithley Instruments Inc., Model 2635B']

    # define measurement paramters here
    averages = IntegerParameter('Averages', default=50)
    measurement_delay = FloatParameter('Measurement Delay', default=0.5)
    max_voltage = FloatParameter('Maximum Voltage', units='V', default=0.5)
    min_voltage = IntegerParameter('Minimum Voltage', units='V', default=-1.0)
    compliance = FloatParameter('Compliance', units='A', default=0.1)
    cycles = IntegerParameter('No. of Cycles', default=1)
    voltage_step = FloatParameter('Voltage Step', units='V', default=0.1)

    # Add Comments as parameters to show up in measurement file
    operator = Parameter('Operator', default='JD')
    location = Parameter('Location', default='Mun')
    setup = Parameter('Setup', default='Probe Station')

    # Calculate the number of data points from range and step
    data_points = IntegerParameter('Data points',
                                   default=np.ceil((max_voltage.value - min_voltage.value) / voltage_step.value))

    # define DATA_COLUMNS that are written to the file
    DATA_COLUMNS = ['Voltage (V)', 'Current (A)']

    def startup(self):
        print('startup')
        for adress,name in self.instruments_dict.items():
            if 'Keithley Instruments Inc., Model 2635B' in name:
                self.instrument_adress = adress
        log.info("Connecting and configuring the instrument")
        log.info("Instrument Adress: " + self.instrument_adress)
        log.info("Instrument Dict: " + str(self.instruments_dict))
        self.sourcemeter = Keithley2600(self.instrument_adress)
        self.sourcemeter.triad()
        self.sourcemeter.set_screentext('$R PulseIVCycle $N$B Ready to measure')

    def execute(self):
        print('execute')
        # reset instrument and its dedicated buffer
        self.sourcemeter.reset()
        self.sourcemeter.clear_buffer()
        self.sourcemeter.setup_buffer(precision=6)
        log.info(f'start: {self.min_voltage}. stop {self.max_voltage}, stime {self.measurement_delay}. points = {self.data_points}')

        self.sourcemeter.set_output(state='ON')
        self.sourcemeter.auto_sweep(start=0, stop=self.max_voltage, stime=self.measurement_delay, points=np.ceil(self.data_points / 2 +1),
                                   source='V')
        self.sourcemeter.wait_for_srq()
        results = self.sourcemeter.get_buffer_data()
        for i in range(0, len(results['sourced']) - 1):
            self.emit('results',
                      {
                          'Voltage (V)': results['sourced'][i],
                          'Current (A)': results['measured'][i],
                      })

    def shutdown(self):
        self.sourcemeter.shutdown()
        log.info("Finished measuring")
        print('shutdown')

class PulseIVCycle(ProcedureWithInstruments):
    """
    Uses a Keithley 26XX device to perform the following measurement:
    1. Pulse at voltage `pulse_voltage` for `pulse_duration` in ms
    2. Perform `cycles` sweeps from `min_voltage` to `max_voltage` starting with a sweep from 0 to max and ending with a sweep from max to 0
    """

    # define measurement paramters here
    max_voltage = FloatParameter('Maximum Voltage', units='V', default=1)
    min_voltage = FloatParameter('Minimum Voltage', units='V',
                                   default=-1)
    compliance = FloatParameter('Compliance', units='A', default=0.1)
    cycles = IntegerParameter('No. of Cycles', default=1)
    voltage_step = FloatParameter('Voltage Step', units='V', default=0.1)
    stime = FloatParameter('Settling Time', units='s', default=0.5)

    # Add Comments as parameters to show up in measurement file
    operator = Parameter('Operator', default='JD')
    location = Parameter('Location', default='Mun')
    setup = Parameter('Setup', default='Probe Station')

    # define DATA_COLUMNS that are written to the file
    DATA_COLUMNS = ['Voltage (V)', 'Current (A)', 'Current Std (A)', 'Cycle']

    def startup(self):
        # System startup: Build instances of all necessary device objects here
        log.info("Connecting and configuring the instrument")
        log.info("Instrument Adress" + instrument_adress)
        self.sourcemeter = Keithley2600(instrument_adress)
        self.sourcemeter.reset()
        self.sourcemeter.clear_buffer()
        self.sourcemeter.triad()
        self.sourcemeter.set_screentext('$R PulseIVCycle $N$B Ready to measure')

    def execute(self):
        # Make Pulse

        # Make Sweep
        log.info("Starting sweep")
        steps = (self.max_voltage - self.min_voltage) / self.voltage_step
        self.sourcemeter.autosweep(0, self.max_voltage, self.stime, steps, 'lin', 'V')
        for i in cycles:
            if self.should_stop():
                log.info("User aborted the procedure")
                break
            log.info(f"Performing sweep number {i}")
            self.soucemeter.autosweep(self.max_voltage, self.min_voltage, self.stime, steps, 'lin', 'V')
            # TODO:Wait for finished measurement?!
            self.soucemeter.autosweep(self.min_voltage, self.max_voltage, self.stime, steps, 'lin', 'V')
            # TODO:Wait for finished measurement?!
        else:
            self.sourcemeter.autosweep(self.max_voltage, 0, self.stime, steps, 'lin', 'V')
            # TODO:Wait for finished measurement?!

        data_array = self.sourcemeter.get_buffer_data()

        # print to console to check array
        print(data_array)
        # emit data
        for i in range(0, len(data_array) - 1):
            self.emit('results',
                      {
                          'Voltage (V)': data_array[0][i],
                          'Current (A)': data_array[1][i],
                          'Voltage Std (V)': data_array[2][i]
                      })

    def shutdown(self):
        self.sourcemeter.shutdown()

class Retention(ProcedureWithInstruments):
    # retention here
    # paramet

    averages = IntegerParameter('Averages', default=50)
    set_voltage = FloatParameter('SET Voltage', units='V', default=1)
    read_voltage = FloatParameter('READ Voltage', units='V', default=1)
    set_duration = FloatParameter('Duration SET Pulse', units='ms', default=1)
    read_duration = FloatParameter('Duration READ Pulse', units='ms', default=1)
    read_delay = FloatParameter('WRITE-READ Delay', units='s', default=1)
    cycles = IntegerParameter('#Cycles', default=1)
    DATA_COLUMNS = ['Voltage (V)', 'Current (A)', 'Current Std (A)', 'Cycle']

    def execute(self):
        log.info("Retention Procedure Running")

    def shutdown(self):
        self.sourcemeter.shutdown()
        log.info("Finished measuring")

class RandomProcedure(ProcedureWithInstruments):
    iterations = IntegerParameter('Loop Iterations')
    delay = FloatParameter('Delay Time', units='s', default=0.2)
    seed = Parameter('Random Seed', default='12345')
    negNumber = IntegerParameter('Number', default=-1)

    DATA_COLUMNS = ['Iteration', 'Random Number']

    def startup(self):
        log.info("Setting the seed of the random number generator")
        log.info("Connected Instruments:"+str(self.instruments_dict))
        random.seed(self.seed)

    def execute(self):
        log.info("Starting the loop of %d iterations" % self.iterations)
        for i in range(self.iterations):
            data = {
                'Iteration': i,
                'Random Number': random.random()
            }
            self.emit('results', data)
            log.debug("Emitting results: %s" % data)
            sleep(self.delay)
            if self.should_stop():
                log.warning("Caught the stop flag in the procedure")
                break


import cv2
import numpy as np
from .grasp_learner import grasp_obj
from .grasp_predictor import Predictors
from .shake_learner import shake_obj
from .shake_predictor import Predictors as Adv_Predictors
import time
import os
import ipdb

from grasp.utils.mjcf_utils import xml_path_completion
from grasp.utils.mjcf_utils import image_path_completion
from grasp.utils.mjcf_utils import root_path_completion
from grasp.utils.mjcf_utils import root_path_completion2
from grasp.utils.mjcf_utils import predict_image_path_completion

from mujoco_py import MjSim, MjViewer, load_model_from_path, functions
from mujoco_py.generated import const

from termcolor import colored

def drawRectangle(I, h, w, t, gsize=300):
    I_temp = I
    grasp_l = gsize/2.5
    grasp_w = gsize/5.0
    grasp_angle = t*(np.pi/18)-np.pi/2

    points = np.array([[-grasp_l, -grasp_w],
                       [grasp_l, -grasp_w],
                       [grasp_l, grasp_w],
                       [-grasp_l, grasp_w]])
    R = np.array([[np.cos(grasp_angle), -np.sin(grasp_angle)],
                  [np.sin(grasp_angle), np.cos(grasp_angle)]])
    rot_points = np.dot(R, points.transpose()).transpose()
    im_points = rot_points + np.array([w,h])
    # print('rec points: ',im_points)

    min_x = min(im_points[0][0], im_points[2][0], im_points[2][0], im_points[3][0])
    max_x = max(im_points[0][0], im_points[1][0], im_points[2][0], im_points[3][0])
    min_y = min(im_points[0][1], im_points[1][1], im_points[2][1], im_points[3][1])
    max_y = max(im_points[0][1], im_points[1][1], im_points[2][1], im_points[3][1])
    center_pt = (int((min_x + max_x) / 2), int((min_y + max_y) / 2))
    print('center points: ', center_pt)
    cv2.circle(I_temp, center_pt, 2, (255, 0, 0), -1)

    cv2.line(I_temp, tuple(im_points[0].astype(int)), tuple(im_points[1].astype(int)), color=(0,255,0), thickness=5)
    cv2.line(I_temp, tuple(im_points[1].astype(int)), tuple(im_points[2].astype(int)), color=(0,0,255), thickness=5)
    cv2.line(I_temp, tuple(im_points[2].astype(int)), tuple(im_points[3].astype(int)), color=(0,255,0), thickness=5)
    cv2.line(I_temp, tuple(im_points[3].astype(int)), tuple(im_points[0].astype(int)), color=(0,0,255), thickness=5)
    return (I_temp, center_pt, R, grasp_angle)

# if train: root_path_completion
# if inference: root_path_completion2
def init_detector(num_samples,use_pro_new='False', use_pro_name='', batch_size=1, gpu_id=0, lr_rate=0.01, test_user= False):
    # init
    if test_user:
        if use_pro_new and os.path.exists(root_path_completion2('predict_module/models/pro_model{}.index'.format(use_pro_name))):
            model_path = root_path_completion2('predict_module/models/pro_model{}'.format(use_pro_name))
        else:
            model_path = root_path_completion2('predict_module/models/Grasp_model')
    else:
        if use_pro_new and os.path.exists(root_path_completion('predict_module/models/pro_model{}.index'.format(use_pro_name))):
            model_path = root_path_completion('predict_module/models/pro_model{}'.format(use_pro_name))
        else:
            model_path = root_path_completion('predict_module/models/Grasp_model')

    print('Loading grasp model')
    G = grasp_obj(model_path, gpu_id, num_samples)
    G.BATCH_SIZE = batch_size
    G.test_init(lr_rate)

    return G

# debug
def init_adv(num_samples, use_new_model='False', use_new_name='', batch_size=1, gpu_id=0,adv_lr_rate=0.01):
    if use_new_model and os.path.exists(root_path_completion('predict_module/models/adv_model{}.index'.format(use_new_name))):
        model_path = root_path_completion('predict_module/models/adv_model{}'.format(use_new_name))
    else:
        model_path = root_path_completion('predict_module/models/Shake_model')

    print('Loading shake model')
    G= shake_obj(model_path, gpu_id, num_samples)
    G.BATCH_SIZE = batch_size
    G.test_init(adv_lr_rate)
    return G

# nbatches is batch size
def predict_from_img(image, obj_name, G, timestep, log_dir, is_train=True):
    gscale = 0.234
    I = cv2.imread(image)
    imsize = 896
    gsize = int(gscale * imsize)

    # detector needs to have been initialized
    P = Predictors(I, G)

    print('Predicting on samples')
    st_time = time.time()

    print(colored('is_train in predict_from_img: {}'.format(is_train),'red'))
    P.graspNet_grasp(patch_size=gsize, num_samples=G.BATCH_SIZE, is_train = is_train)
    #P.graspNet_grasp(patch_size=gsize, num_samples=1000)
    print('Time taken: {}s'.format(time.time() - st_time))

    print('debug predict_from_img fc8_norm_vals shape: ', P.fc8_norm_vals.shape)

    if not is_train:
        index = np.argmax(P.fc8_norm_vals)
        r_pindex = index // P.fc8_norm_vals.shape[1]
        r_tindex = index%P.fc8_norm_vals.shape[1]
    else:
        r_pindex = P.r_pindex
        r_tindex = P.r_tindex

    (I, center_pt, R, angle) = drawRectangle(I, P.patch_hs[r_pindex], P.patch_ws[r_pindex], r_tindex, gsize)

    patch_Is_resized = P.patch_Is_resized[r_pindex]
    patch_Is_resized = patch_Is_resized[np.newaxis, :] #[1,224,224,3]

    # save_name = image_path_completion(obj_name + '_predict.jpg')
    save_name = predict_image_path_completion(obj_name + '_predict_{}.jpg'.format(timestep), log_dir)
    try:
        cv2.imwrite(save_name, I)
        print('prediction image {} saved'.format(save_name))
    except:
        print('error saving image: {}'.format(save_name))

    # G.test_close()
    pred = P.pred
    pred = np.max(pred) #[1]
    del P
    # one patch, one prediction [1, 18], one groundtruth [1,18]
    return (center_pt, R, angle, patch_Is_resized, pred)

# debug, adv_precition using shakeNet
def adv_predict_from_img(I, G_adv):
    A = Adv_Predictors(I, G_adv)
    print('Adv Predicting on samples')
    st_time = time.time()
    probs, adv_action = A.shakeNet_shake(num_actions= 6, num_samples = 1)
    print('Time take: {}s'.format(time.time()-st_time))
    del A

    return probs,adv_action