import os, math, ctypes
import numpy as np
from OpenGL import GL
from art import VERT_LENGTH
from palette import MAX_COLORS

# inactive layer alphas
LAYER_VIS_FULL = 1
LAYER_VIS_DIM = 0.25
LAYER_VIS_NONE = 0

class TileRenderable:
    """
    3D visual representation of an Art. Each layer is rendered as grids of
    rectangular OpenGL triangle-pairs. Animation frames are uploaded into our
    buffers from source Art's numpy arrays.
    """
    vert_shader_source = 'renderable_v.glsl'
    "vertex shader: includes view projection matrix, XYZ camera uniforms."
    frag_shader_source = 'renderable_f.glsl'
    "Pixel shader: handles FG/BG colors."
    log_create_destroy = False
    log_animation = False
    log_buffer_updates = False
    grain_strength = 0.
    alpha = 1.
    "Alpha (0 to 1) for entire Renderable."
    bg_alpha = 1.
    "Alpha (0 to 1) *only* for tile background colors."
    default_move_rate = 1
    use_art_offset = True
    "Use game object's art_off_pct values."

    def __init__(self, app, art, game_object=None):
        "Create Renderable with given Art, optionally bound to given GameObject"
        self.app = app
        "Application that renders us."
        self.art = art
        "Art we get data from."
        self.art.renderables.append(self)
        self.go = game_object
        "GameObject we're attached to."
        self.exporting = False
        "Set True momentarily by image export process; users shouldn't touch."
        self.visible = True
        "Boolean for easy render / don't-render functionality."
        self.frame = self.art.active_frame or 0
        "Frame of our art's animation we're currently on"
        self.animating = False
        self.anim_timer, self.last_frame_time = 0, 0
        # world space position and scale
        self.x, self.y, self.z = 0, 0, 0
        self.scale_x, self.scale_y, self.scale_z = 1, 1, 1
        if self.go:
            self.scale_x = self.go.scale_x
            self.scale_y = self.go.scale_y
            self.scale_z = self.go.scale_z
        # width and height in XY render space
        self.width, self.height = 1, 1
        self.reset_size()
        # TODO: object rotation matrix, if needed
        self.goal_x, self.goal_y, self.goal_z = 0, 0, 0
        self.move_rate = self.default_move_rate
        # marked True when UI is interpolating it
        self.ui_moving = False
        self.camera = self.app.camera
        # bind VAO etc before doing shaders etc
        if self.app.use_vao:
            self.vao = GL.glGenVertexArrays(1)
            GL.glBindVertexArray(self.vao)
        self.shader = self.app.sl.new_shader(self.vert_shader_source, self.frag_shader_source)
        self.proj_matrix_uniform = self.shader.get_uniform_location('projection')
        self.view_matrix_uniform = self.shader.get_uniform_location('view')
        self.position_uniform = self.shader.get_uniform_location('objectPosition')
        self.scale_uniform = self.shader.get_uniform_location('objectScale')
        self.charset_width_uniform = self.shader.get_uniform_location('charMapWidth')
        self.charset_height_uniform = self.shader.get_uniform_location('charMapHeight')
        self.char_uv_width_uniform = self.shader.get_uniform_location('charUVWidth')
        self.char_uv_height_uniform = self.shader.get_uniform_location('charUVHeight')
        self.charset_tex_uniform = self.shader.get_uniform_location('charset')
        self.palette_tex_uniform = self.shader.get_uniform_location('palette')
        self.grain_tex_uniform = self.shader.get_uniform_location('grain')
        self.palette_width_uniform = self.shader.get_uniform_location('palTextureWidth')
        self.grain_strength_uniform = self.shader.get_uniform_location('grainStrength')
        self.alpha_uniform = self.shader.get_uniform_location('alpha')
        self.brightness_uniform = self.shader.get_uniform_location('brightness')
        self.bg_alpha_uniform = self.shader.get_uniform_location('bgColorAlpha')
        self.create_buffers()
        # finish
        if self.app.use_vao:
            GL.glBindVertexArray(0)
        if self.log_create_destroy:
            self.app.log('created: %s' % self)

    def __str__(self):
        "for debug purposes, return a concise unique name"
        for i,r in enumerate(self.art.renderables):
            if r is self:
                break
        return '%s %s %s' % (self.art.get_simple_name(), self.__class__.__name__, i)

    def create_buffers(self):
        # vertex positions and elements
        # determine vertex count needed for render
        self.vert_count = int(len(self.art.elem_array))
        self.vert_buffer, self.elem_buffer = GL.glGenBuffers(2)
        self.update_buffer(self.vert_buffer, self.art.vert_array,
                           GL.GL_ARRAY_BUFFER, GL.GL_STATIC_DRAW, GL.GL_FLOAT, 'vertPosition', VERT_LENGTH)
        self.update_buffer(self.elem_buffer, self.art.elem_array,
                           GL.GL_ELEMENT_ARRAY_BUFFER, GL.GL_STATIC_DRAW, GL.GL_UNSIGNED_INT, None, None)
        # tile data buffers
        # use GL_DYNAMIC_DRAW given they change every time a char/color changes
        self.char_buffer, self.uv_buffer = GL.glGenBuffers(2)
        # character indices (which become vertex UVs)
        self.update_buffer(self.char_buffer, self.art.chars[self.frame],
                           GL.GL_ARRAY_BUFFER, GL.GL_DYNAMIC_DRAW, GL.GL_FLOAT, 'charIndex', 1)
        # UV "mods" - modify UV derived from character index
        self.update_buffer(self.uv_buffer, self.art.uv_mods[self.frame],
                           GL.GL_ARRAY_BUFFER, GL.GL_DYNAMIC_DRAW, GL.GL_FLOAT, 'uvMod', 2)
        self.fg_buffer, self.bg_buffer = GL.glGenBuffers(2)
        # foreground/background color indices (which become rgba colors)
        self.update_buffer(self.fg_buffer, self.art.fg_colors[self.frame],
                           GL.GL_ARRAY_BUFFER, GL.GL_DYNAMIC_DRAW, GL.GL_FLOAT, 'fgColorIndex', 1)
        self.update_buffer(self.bg_buffer, self.art.bg_colors[self.frame],
                           GL.GL_ARRAY_BUFFER, GL.GL_DYNAMIC_DRAW, GL.GL_FLOAT, 'bgColorIndex', 1)

    def update_geo_buffers(self):
        self.update_buffer(self.vert_buffer, self.art.vert_array, GL.GL_ARRAY_BUFFER, GL.GL_STATIC_DRAW, GL.GL_FLOAT, None, None)
        self.update_buffer(self.elem_buffer, self.art.elem_array, GL.GL_ELEMENT_ARRAY_BUFFER, GL.GL_STATIC_DRAW, GL.GL_UNSIGNED_INT, None, None)
        # total vertex count probably changed
        self.vert_count = int(len(self.art.elem_array))

    def update_tile_buffers(self, update_chars, update_uvs, update_fg, update_bg):
        "Update GL data arrays for tile characters, fg/bg colors, transforms."
        updates = {}
        if update_chars:
            updates[self.char_buffer] = self.art.chars
        if update_uvs:
            updates[self.uv_buffer] = self.art.uv_mods
        if update_fg:
            updates[self.fg_buffer] = self.art.fg_colors
        if update_bg:
            updates[self.bg_buffer] = self.art.bg_colors
        for update in updates:
            self.update_buffer(update, updates[update][self.frame],
                               GL.GL_ARRAY_BUFFER, GL.GL_DYNAMIC_DRAW,
                               GL.GL_FLOAT, None, None)

    def update_buffer(self, buffer_index, array, target, buffer_type, data_type,
                      attrib_name, attrib_size):
        if self.log_buffer_updates:
            self.app.log('update_buffer: %s, %s, %s, %s, %s, %s, %s' % (buffer_index, array, target, buffer_type, data_type, attrib_name, attrib_size))
        GL.glBindBuffer(target, buffer_index)
        GL.glBufferData(target, array.nbytes, array, buffer_type)
        if attrib_name:
            attrib = self.shader.get_attrib_location(attrib_name)
            GL.glEnableVertexAttribArray(attrib)
            GL.glVertexAttribPointer(attrib, attrib_size, data_type,
                                     GL.GL_FALSE, 0, ctypes.c_void_p(0))
        # unbind each buffer before binding next
        GL.glBindBuffer(target, 0)

    def advance_frame(self):
        "Advance to our Art's next animation frame."
        self.set_frame(self.frame + 1)

    def rewind_frame(self):
        "Rewind to our Art's previous animation frame."
        self.set_frame(self.frame - 1)

    def set_frame(self, new_frame_index):
        "Set us to display our Art's given animation frame."
        if new_frame_index == self.frame:
            return
        old_frame = self.frame
        self.frame = new_frame_index % self.art.frames
        self.update_tile_buffers(True, True, True, True)
        if self.log_animation:
            self.app.log('%s animating from frames %s to %s' % (self, old_frame, self.frame))

    def start_animating(self):
        "Start animation playback."
        self.animating = True
        self.anim_timer = 0

    def stop_animating(self):
        "Pause animation playback on current frame (in game mode)."
        self.animating = False
        # restore to active frame if stopping
        if not self.app.game_mode:
            self.set_frame(self.art.active_frame)

    def set_art(self, new_art):
        "Display and bind to given Art."
        if self.art:
            self.art.renderables.remove(self)
        self.art = new_art
        self.reset_size()
        self.art.renderables.append(self)
        # make sure frame is valid
        self.frame %= self.art.frames
        self.update_geo_buffers()
        self.update_tile_buffers(True, True, True, True)
        #print('%s now uses Art %s' % (self, self.art.filename))

    def reset_size(self):
        self.width = self.art.width * self.art.quad_width * abs(self.scale_x)
        self.height = self.art.height * self.art.quad_height * self.scale_y

    def move_to(self, x, y, z, travel_time=None):
        """
        Start simple linear interpolation to given destination over given time.
        Not very useful in Game Mode, mainly used for document switch UI effect
        in Art Mode.
        """
        # for fixed travel time, set move rate accordingly
        if travel_time:
            frames = (travel_time * 1000) / max(self.app.framerate, 30)
            dx = x - self.x
            dy = y - self.y
            dz = z - self.z
            dist = math.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
            self.move_rate = dist / frames
        else:
            self.move_rate = self.default_move_rate
        self.ui_moving = True
        self.goal_x, self.goal_y, self.goal_z = x, y, z
        if self.log_animation:
            self.app.log('%s will move to %s,%s' % (self.art.filename, self.goal_x, self.goal_y))

    def snap_to(self, x, y, z):
        self.x, self.y, self.z = x, y, z
        self.goal_x, self.goal_y, self.goal_z = x, y, z
        self.ui_moving = False

    def update_transform_from_object(self, obj):
        "Update our position & scale based on that of given game object."
        self.z = obj.z
        if self.scale_x != obj.scale_x or self.scale_y != obj.scale_y:
            self.reset_size()
        self.x, self.y = obj.x, obj.y
        if self.use_art_offset:
            if obj.flip_x:
                self.x += self.width * obj.art_off_pct_x
            else:
                self.x -= self.width * obj.art_off_pct_x
            self.y += self.height * obj.art_off_pct_y
        self.scale_x, self.scale_y = obj.scale_x, obj.scale_y
        if obj.flip_x:
            self.scale_x *= -1
        self.scale_z = obj.scale_z

    def update_loc(self):
        # TODO: probably time to bust out the ol' vector module for this stuff
        # get delta
        dx = self.goal_x - self.x
        dy = self.goal_y - self.y
        dz = self.goal_z - self.z
        dist = math.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
        # close enough?
        if dist <= self.move_rate:
            self.x = self.goal_x
            self.y = self.goal_y
            self.z = self.goal_z
            self.ui_moving = False
            return
        # normalize
        inv_dist = 1 / dist
        dir_x = dx * inv_dist
        dir_y = dy * inv_dist
        dir_z = dz * inv_dist
        self.x += self.move_rate * dir_x
        self.y += self.move_rate * dir_y
        self.z += self.move_rate * dir_z
        #self.app.log('%s moved to %s,%s' % (self, self.x, self.y))

    def update(self):
        if self.go:
            self.update_transform_from_object(self.go)
        if self.ui_moving:
            self.update_loc()
        if not self.animating:
            return
        if self.app.game_mode and self.app.gw.paused:
            return
        elapsed = self.app.get_elapsed_time() - self.last_frame_time
        self.anim_timer += elapsed
        new_frame = self.frame
        this_frame_delay = self.art.frame_delays[new_frame] * 1000
        while self.anim_timer >= this_frame_delay:
            self.anim_timer -= this_frame_delay
            # iterate through frames, but don't call set_frame until we're done
            new_frame += 1
            new_frame %= self.art.frames
            this_frame_delay = self.art.frame_delays[new_frame] * 1000
            # TODO: if new_frame < self.frame, count anim loop?
        self.set_frame(new_frame)
        self.last_frame_time = self.app.get_elapsed_time()

    def destroy(self):
        if self.app.use_vao:
            GL.glDeleteVertexArrays(1, [self.vao])
        GL.glDeleteBuffers(6, [self.vert_buffer, self.elem_buffer, self.char_buffer, self.uv_buffer, self.fg_buffer, self.bg_buffer])
        if self.art and self in self.art.renderables:
            self.art.renderables.remove(self)
        if self.log_create_destroy:
            self.app.log('destroyed: %s' % self)

    def get_projection_matrix(self):
        """
        UIRenderable overrides this so it doesn't have to override
        Renderable.render and duplicate lots of code.
        """
        return np.eye(4, 4) if self.exporting else self.camera.projection_matrix

    def get_view_matrix(self):
        return np.eye(4, 4) if self.exporting else self.camera.view_matrix

    def get_loc(self):
        "Returns world space location as (x, y, z) tuple."
        export_loc = (-1, 1, 0)
        return export_loc if self.exporting else (self.x, self.y, self.z)

    def get_scale(self):
        "Returns world space scale as (x, y, z) tuple."
        if not self.exporting:
            return (self.scale_x, self.scale_y, self.scale_z)
        x = 2 / (self.art.width * self.art.quad_width)
        y = 2 / (self.art.height * self.art.quad_height)
        return (x, y, 1)

    def render_frame_for_export(self, frame):
        self.exporting = True
        self.set_frame(frame)
        # cache "inactive layer visibility", restore after render
        ilv = self.art.app.inactive_layer_visibility
        self.art.app.inactive_layer_visibility = LAYER_VIS_FULL
        # cursor might be hovering, undo any preview changes
        for edit in self.art.app.cursor.preview_edits:
            edit.undo()
        # update art to commit changes to the renderable
        self.art.update()
        self.render()
        self.art.app.inactive_layer_visibility = ilv
        self.exporting = False

    def render(self, layers=None, z_override=None, brightness=1.0):
        """
        Render given list of layers at given Z depth.
        If layers is None, render all layers.
        If layers is an int, just render that layer.
        If z_override is None, render each layer at Z defined in our Art.
        """
        if not self.visible:
            return
        GL.glUseProgram(self.shader.program)
        # bind textures - character set, palette, UI grain
        GL.glActiveTexture(GL.GL_TEXTURE0)
        GL.glUniform1i(self.charset_tex_uniform, 0)
        GL.glBindTexture(GL.GL_TEXTURE_2D, self.art.charset.texture.gltex)
        GL.glActiveTexture(GL.GL_TEXTURE1)
        GL.glUniform1i(self.palette_tex_uniform, 1)
        GL.glBindTexture(GL.GL_TEXTURE_2D, self.art.palette.texture.gltex)
        GL.glActiveTexture(GL.GL_TEXTURE2)
        GL.glUniform1i(self.grain_tex_uniform, 2)
        GL.glBindTexture(GL.GL_TEXTURE_2D, self.app.ui.grain_texture.gltex)
        # set active texture unit back after binding 2nd-Nth textures
        GL.glActiveTexture(GL.GL_TEXTURE0)
        GL.glUniform1i(self.charset_width_uniform, self.art.charset.map_width)
        GL.glUniform1i(self.charset_height_uniform, self.art.charset.map_height)
        GL.glUniform1f(self.char_uv_width_uniform, self.art.charset.u_width)
        GL.glUniform1f(self.char_uv_height_uniform, self.art.charset.v_height)
        GL.glUniform1f(self.palette_width_uniform, MAX_COLORS)
        GL.glUniform1f(self.grain_strength_uniform, self.grain_strength)
        # camera uniforms
        GL.glUniformMatrix4fv(self.proj_matrix_uniform, 1, GL.GL_FALSE,
                              self.get_projection_matrix())
        GL.glUniformMatrix4fv(self.view_matrix_uniform, 1, GL.GL_FALSE,
                              self.get_view_matrix())
        # TODO: determine if cost of setting all above uniforms for each
        # Renderable is significant enough to warrant opti where they're set once
        GL.glUniform1f(self.bg_alpha_uniform, self.bg_alpha)
        GL.glUniform1f(self.brightness_uniform, brightness)
        GL.glUniform3f(self.scale_uniform, *self.get_scale())
        # VAO vs non-VAO paths
        if self.app.use_vao:
            GL.glBindVertexArray(self.vao)
        else:
            attrib = self.shader.get_attrib_location # for brevity
            vp = ctypes.c_void_p(0)
            # bind each buffer and set its attrib:
            # verts
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vert_buffer)
            GL.glVertexAttribPointer(attrib('vertPosition'), VERT_LENGTH, GL.GL_FLOAT, GL.GL_FALSE, 0, vp)
            GL.glEnableVertexAttribArray(attrib('vertPosition'))
            # chars
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.char_buffer)
            GL.glVertexAttribPointer(attrib('charIndex'), 1, GL.GL_FLOAT, GL.GL_FALSE, 0, vp)
            GL.glEnableVertexAttribArray(attrib('charIndex'))
            # uvs
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.uv_buffer)
            GL.glVertexAttribPointer(attrib('uvMod'), 2, GL.GL_FLOAT, GL.GL_FALSE, 0, vp)
            GL.glEnableVertexAttribArray(attrib('uvMod'))
            # fg colors
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.fg_buffer)
            GL.glVertexAttribPointer(attrib('fgColorIndex'), 1, GL.GL_FLOAT, GL.GL_FALSE, 0, vp)
            GL.glEnableVertexAttribArray(attrib('fgColorIndex'))
            # bg colors
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.bg_buffer)
            GL.glVertexAttribPointer(attrib('bgColorIndex'), 1, GL.GL_FLOAT, GL.GL_FALSE, 0, vp)
            GL.glEnableVertexAttribArray(attrib('bgColorIndex'))
        # finally, bind element buffer
        GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, self.elem_buffer)
        GL.glEnable(GL.GL_BLEND)
        GL.glBlendFunc(GL.GL_SRC_ALPHA, GL.GL_ONE_MINUS_SRC_ALPHA)
        # draw all specified layers if no list given
        if layers is None:
            # sort layers in Z depth
            layers = list(range(self.art.layers))
            layers.sort(key=lambda i: self.art.layers_z[i], reverse=False)
        # handle a single int param
        elif type(layers) is int:
            layers = [layers]
        layer_size = int(len(self.art.elem_array) / self.art.layers)
        for i in layers:
            # skip game mode-hidden layers
            if not self.app.show_hidden_layers and not self.art.layers_visibility[i]:
                continue
            layer_start = i * layer_size
            layer_end = layer_start + layer_size
            # for active art, dim all but active layer based on UI setting
            if not self.app.game_mode and self.art is self.app.ui.active_art and i != self.art.active_layer:
                GL.glUniform1f(self.alpha_uniform, self.alpha * self.app.inactive_layer_visibility)
            else:
                GL.glUniform1f(self.alpha_uniform, self.alpha)
            # use position offset instead of baked-in Z for layers - this
            # way a layer's Z can change w/o rebuilding its vert array
            x, y, z = self.get_loc()
            # for export, render all layers at same Z
            if not self.exporting:
                z += self.art.layers_z[i]
                z = z_override if z_override else z
            GL.glUniform3f(self.position_uniform, x, y, z)
            GL.glDrawElements(GL.GL_TRIANGLES, layer_size, GL.GL_UNSIGNED_INT,
                ctypes.c_void_p(layer_start * ctypes.sizeof(ctypes.c_uint)))
        GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, 0)
        GL.glDisable(GL.GL_BLEND)
        if self.app.use_vao:
            GL.glBindVertexArray(0)
        GL.glUseProgram(0)

class OnionTileRenderable(TileRenderable):

    "TileRenderable subclass used for onion skin display in Art Mode animation."

    # never animate
    def start_animating(self):
        pass

    def stop_animating(self):
        pass

class GameObjectRenderable(TileRenderable):

    """
    TileRenderable subclass used by GameObjects. Almost no custom logic for now.
    """

    def get_loc(self):
        """
        Returns world space location as (x, y, z) tuple, offset by our
        GameObject's location.
        """
        x, y, z = self.x, self.y, self.z
        if self.go:
            off_x, off_y, off_z = self.go.get_render_offset()
            x += off_x
            y += off_y
            z += off_z
        return x, y, z


import numpy
import os
import sys

from nematus.nmt import train
from sequence_qe.dataset import mkdir_p

# the size of the src+mt concatenated data vocab
SRC_AND_MT_VOCAB_SIZE = 38559
FACTOR_1_SIZE = 308
FACTOR_2_SIZE = 508
FACTOR_3_SIZE = 308

PE_VOCAB_SIZE = 41101
# never early stop
PATIENCE = 99999999999

INPUT_VOCAB_DIR = "/media/1tb_drive/Dropbox/data/qe/amunmt_artificial_ape_2016/data/4M/spacy_factor_corpus"
OUTPUT_VOCAB_DIR = "/media/1tb_drive/nematus_ape_experiments/amunmt_ape_pretrained/system/models"
SRC_VOCAB = os.path.join(INPUT_VOCAB_DIR, 'train.src-mt.concatenated.bpe.json')
FACTOR_1_VOCAB = os.path.join(INPUT_VOCAB_DIR, 'factor_1.json')
FACTOR_2_VOCAB = os.path.join(INPUT_VOCAB_DIR, 'factor_2.json')
FACTOR_3_VOCAB = os.path.join(INPUT_VOCAB_DIR, 'factor_3.json')
PE_VOCAB = os.path.join(OUTPUT_VOCAB_DIR, 'mt-pe/vocab.pe.json')

TRAIN_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/amunmt_artificial_ape_2016/data/4M/spacy_factor_corpus"
SRC_TRAIN = os.path.join(TRAIN_DATA_DIR, 'train.src-mt.concatenated.bpe.factor_corpus')
TRG_TRAIN = os.path.join(TRAIN_DATA_DIR, 'train.pe.prepped.orig')

# WMT 16 EN-DE QE/APE DEV Data
QE_DATA_DIR = "/media/1tb_drive/Dropbox/data/qe/ape/concat_wmt_2016_2017/spacy_factor_corpus"
SRC_DEV = os.path.join(QE_DATA_DIR, 'dev.src-mt.concatenated.bpe.factor_corpus')
TRG_DEV = os.path.join(QE_DATA_DIR, 'dev.pe.prepped.orig')

if __name__ == '__main__':
    validerr = train(saveto='model/model.npz',
                    reload_=True,
                    dim_word=512,
                    dim=1028,
                    n_words=PE_VOCAB_SIZE,
                    n_words_src=SRC_AND_MT_VOCAB_SIZE,
                    decay_c=0.,
                    clip_c=1.,
                    lrate=0.0001,
                    optimizer='adadelta',
                    maxlen=100,
                    batch_size=32,
                    valid_batch_size=32,
                    datasets=[SRC_TRAIN, TRG_TRAIN],
                    valid_datasets=[SRC_DEV, TRG_DEV],
                    dictionaries=[SRC_VOCAB, FACTOR_1_VOCAB, FACTOR_2_VOCAB, FACTOR_3_VOCAB, PE_VOCAB],
                    factors=4,
                    dim_per_factor=[256, 64, 128, 64],
                    validFreq=5000,
                    dispFreq=500,
                    saveFreq=5000,
                    sampleFreq=5000,
                    overwrite=False,
                    external_validation_script='./validate.sh')
    print validerr


import json
import logging
import math
import os
import sys
from io import open
from typing import Callable, List, Tuple
import numpy as np
import copy

import torch
from torch import nn
from torch import Tensor, device, dtype

from transformers import BertPreTrainedModel

from .vision_transformer import vit_base_patch16_224

from .vilmodel import BertEmbeddings, ImageEmbeddings, HistoryEmbeddings, LxmertEncoder

class NavTHORImagePreTrainedModel(BertPreTrainedModel):
    r""" Modification of LXMERT Model """
    def __init__(self, config):
        super().__init__(config)
        self.vision_backbone = vit_base_patch16_224(pretrained=True,
            drop_rate=config.hidden_dropout_prob,
            attn_drop_rate=config.attention_probs_dropout_prob,
            drop_path_rate=0.)

        self.embeddings = BertEmbeddings(config)
        self.img_embeddings = ImageEmbeddings(config)
        # share image encoding
        self.hist_embeddings = HistoryEmbeddings(config)

        self.encoder = LxmertEncoder(config)

        self.init_weights()

    def forward_vision_backbone(self, images, detach=False):
        # due to memory issue, we cannot propagate to pano images in the history
        is_pano = len(images.size()) == 6
        if is_pano:
            N, T, P, C, H, W = images.size()    # pano images
            images = images.view(N*T*P, C, H, W)
        else:
            N, T, C, H, W = images.size()
            images = images.view(N*T, C, H, W)
        # feats = self.vision_backbone.forward_features(images)
        if is_pano:
            with torch.no_grad():
                feats = self.vision_backbone.forward_features(images)
            feats = feats.view(N, T, P, -1)
        else:
            feats = self.vision_backbone.forward_features(images)
            feats = feats.view(N, T, -1)
        if detach:
            feats = feats.detach()
        return feats

    def forward(self, txt_ids, txt_masks,
                hist_images, hist_ang_feats, hist_pano_images, hist_pano_ang_feats, hist_masks,
                ob_images, ob_ang_feats, ob_nav_types, ob_masks,
                hist_mrc_masks=None, ob_v_exists=None):
        batch_size = txt_ids.size(0)

        # text embedding
        extended_txt_masks = txt_masks.unsqueeze(1).unsqueeze(2)
        extended_txt_masks = extended_txt_masks.to(dtype=self.dtype)
        extended_txt_masks = (1.0 - extended_txt_masks) * -10000.0
        txt_token_type_ids = torch.zeros_like(txt_ids)
        txt_embeds = self.embeddings(txt_ids, token_type_ids=txt_token_type_ids)

        # history embedding
        extended_hist_masks = hist_masks.unsqueeze(1).unsqueeze(2)
        extended_hist_masks = extended_hist_masks.to(dtype=self.dtype)
        extended_hist_masks = (1.0 - extended_hist_masks) * -10000.0
        if hist_images is not None:
            hist_max_len = hist_images.size(1)
            hist_img_feats = self.forward_vision_backbone(hist_images)
            hist_pano_img_feats = self.forward_vision_backbone(hist_pano_images, detach=True)
            hist_step_ids = torch.arange(hist_max_len).expand((1, -1)).to(self.device)
            if hist_mrc_masks is not None:  # (N, T)
                hist_img_feats = hist_img_feats.masked_fill(hist_mrc_masks.unsqueeze(-1), 0)
                hist_pano_img_feats = hist_pano_img_feats.masked_fill(hist_mrc_masks.unsqueeze(-1).unsqueeze(-1), 0)
        else:
            hist_step_ids = None
            hist_img_feats = None
            hist_pano_img_feats = None
        hist_cls_embeds, hist_vp_embeds = self.hist_embeddings(hist_img_feats, hist_ang_feats,
            hist_pano_img_feats, hist_pano_ang_feats, hist_step_ids,
            batch_size=batch_size)
        if hist_vp_embeds is None:
            hist_embeds = hist_cls_embeds
        else:
            hist_embeds = torch.cat([hist_cls_embeds, hist_vp_embeds], dim=1)

        # image embedding
        if ob_images is not None:
            ob_img_feats = self.forward_vision_backbone(ob_images)
            if ob_v_exists is not None: # [N, ]
                ob_img_feats = ob_img_feats.masked_fill(ob_v_exists.logical_not().unsqueeze(-1).unsqueeze(-1), 0)
            # add STOP token
            ob_img_feats= torch.cat([ob_img_feats, \
                torch.zeros(batch_size, 1, ob_img_feats.size(2), dtype=ob_img_feats.dtype, device=ob_img_feats.device)],
                dim=1)
            ob_token_type_ids = torch.ones(batch_size, 1).long().to(self.device)
            ob_embeds = self.img_embeddings(ob_img_feats, ob_ang_feats,
                self.embeddings.token_type_embeddings(ob_token_type_ids),
                nav_types=ob_nav_types)
            extended_ob_masks = ob_masks.unsqueeze(1).unsqueeze(2)
            extended_ob_masks = extended_ob_masks.to(dtype=self.dtype)
            extended_ob_masks = (1.0 - extended_ob_masks) * -10000.0
        else:
            ob_embeds, extended_ob_masks = None, None

        # multi-modal encoding
        txt_embeds, hist_embeds, ob_embeds = self.encoder(
            txt_embeds, extended_txt_masks,
            hist_embeds, extended_hist_masks,
            ob_embeds, extended_ob_masks)

        return txt_embeds, hist_embeds, ob_embeds

    def forward_itm(self, txt_ids, txt_masks,
                hist_images, hist_ang_feats, hist_pano_images, hist_pano_ang_feats, hist_masks,
                num_neg_trajs=4):
        batch_size, hist_max_len = hist_images.size()[:2]

        # text encoding
        extended_txt_masks = txt_masks.unsqueeze(1).unsqueeze(2)
        extended_txt_masks = extended_txt_masks.to(dtype=self.dtype)
        extended_txt_masks = (1.0 - extended_txt_masks) * -10000.0
        txt_token_type_ids = torch.zeros_like(txt_ids)
        txt_embeds = self.embeddings(txt_ids, token_type_ids=txt_token_type_ids)
        for layer_module in self.encoder.layer:
            temp_output = layer_module(txt_embeds, extended_txt_masks)
            txt_embeds = temp_output[0]
        # copy txt_embeds
        batch_size = txt_embeds.size(0)
        txt_embeds = txt_embeds.repeat(1+num_neg_trajs, 1, 1)
        extended_txt_masks = extended_txt_masks.repeat(1+num_neg_trajs, 1, 1, 1)

        # history encoding
        extended_hist_masks = hist_masks.unsqueeze(1).unsqueeze(2)
        extended_hist_masks = extended_hist_masks.to(dtype=self.dtype)
        extended_hist_masks = (1.0 - extended_hist_masks) * -10000.0

        hist_img_feats = self.forward_vision_backbone(hist_images)
        hist_pano_img_feats = self.forward_vision_backbone(hist_pano_images, detach=True)

        hist_cls_embeds, hist_vp_embeds_no_pos = self.hist_embeddings(hist_img_feats, hist_ang_feats,
            hist_pano_img_feats, hist_pano_ang_feats, pos_ids=None,
            batch_size=batch_size)
        hist_step_ids = torch.arange(hist_max_len).expand((1, -1)).to(self.device)
        hist_vp_embeds = self.hist_embeddings.dropout(self.hist_embeddings.layer_norm(
            hist_vp_embeds_no_pos + self.hist_embeddings.position_embeddings(hist_step_ids)))
        hist_embeds = torch.cat([hist_cls_embeds, hist_vp_embeds], dim=1)

        if self.encoder.h_layers is not None:
            for layer_module in self.encoder.h_layers:
                temp_output = layer_module(hist_embeds, extended_hist_masks)
                hist_embeds = temp_output[0]

        # random negs in batch
        K = num_neg_trajs // 2
        neg_idxs = []
        for i in range(batch_size):
            neg_idxs.append(np.random.choice(np.arange(0, i).tolist() + np.arange(i+1, batch_size).tolist(), K))
        neg_idxs = torch.from_numpy(np.stack(neg_idxs, 0)).to(self.device)

        neg_hist_embeds, neg_hist_masks = [], []
        for k in range(K):
            neg_hist_embeds.append(hist_embeds[neg_idxs[:, k]])
            neg_hist_masks.append(extended_hist_masks[neg_idxs[:, k]])

        # shuffled negs
        hist_lens = torch.sum(hist_masks, 1) - 1
        for _ in range(K):
            shuffled_pos_ids = []
            for i in range(batch_size):
                shuffled_idxs = torch.randperm(hist_lens[i])
                shuffled_idxs = torch.cat([shuffled_idxs, torch.arange(hist_lens[i], hist_max_len, dtype=torch.long)], 0).to(self.device)
                shuffled_pos_ids.append(shuffled_idxs)
            shuffled_pos_ids = torch.stack(shuffled_pos_ids, 0)
            shuffled_hist_embeds = torch.cat([hist_cls_embeds, \
                self.hist_embeddings.dropout(self.hist_embeddings.layer_norm(
                hist_vp_embeds_no_pos + self.hist_embeddings.position_embeddings(shuffled_pos_ids)))], dim=1)

            if self.encoder.h_layers is not None:
                for layer_module in self.encoder.h_layers:
                    temp_output = layer_module(shuffled_hist_embeds, extended_hist_masks)
                    shuffled_hist_embeds = temp_output[0]
            neg_hist_embeds.append(shuffled_hist_embeds)
            neg_hist_masks.append(extended_hist_masks)

        pos_neg_hist_embeds = torch.cat([hist_embeds] + neg_hist_embeds, 0)
        pos_neg_hist_masks = torch.cat([extended_hist_masks] + neg_hist_masks, 0)

        # multi-modal encoding
        for layer_module in self.encoder.x_layers:
            txt_embeds, pos_neg_hist_embeds = layer_module(
                txt_embeds, extended_txt_masks,
                pos_neg_hist_embeds, pos_neg_hist_masks)

        fused_embeds = txt_embeds[:, 0] * pos_neg_hist_embeds[:, 0]
        fused_embeds = torch.stack(torch.split(fused_embeds, batch_size), 1)
        return fused_embeds


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import pandas as pd
from nltk.tokenize import word_tokenize
import numpy as np
import functions as fn

class Rules:
    def __init__(self,Positive,Negative,annot_df,covered,abstract,query_gene,query_id,sentence):
        # loading lookup files
        # define rank variable
#        self.posv=posv
#        self.negn=negn
#        self.posn=posn
#        self.negv=negv
#        self.neutn=neutn
#        self.neutv=neutv
        self.Positive=Positive
        self.Negative=Negative
        self.annot_df=annot_df #list of all annotations for the abstract
        self.covered=covered #list of annotations covered by sentences including query genes
        self.abstract=abstract
        self.query_id=query_id
        self.query_gene=query_gene
        self.sentence=sentence
        self.rank=0
        self.tfrank=0
        self.tgrank=0
        self.direction=0 # +1 : TF - Target // -1 : Target - TF // 0 : not clear

    def search_ranking(self):
        word_tokens = word_tokenize(self.sentence.sentence)
        word_tokens = [w for w in word_tokens if (len(w)>1)]
        covered=self.sentence.covered
        query_list=self.annot_df.loc[covered,['id','text','offset']]
        ids=self.query_id[0]+self.query_id[1]
        query_list=query_list.loc[query_list["id"].isin(ids),:]
        target_list=query_list.loc[query_list["id"].isin(self.query_id[1]),:]
        voting_list=[]
        rank=0
        for target in target_list.itertuples():
            pos_event=0
            neg_event=0
            class_list=[]
            for item in word_tokens:
                try:
                    offset=self.sentence.start+self.sentence.sentence.index(item) #find the offset of item
                    if(item in self.Positive):
                        class_list.append([offset,1])
                        pos_event+=1
                    elif(item in self.Negative):
                        class_list.append([offset,-1])
                        neg_event+=1
                except:
                    continue
            if (pos_event==0 and neg_event==0):
                rank+=0 # when no positive/negative events found in sentence
                break
            if(pos_event==len(target_list.index) and neg_event==0):
                rank+=10 #when all events are found in positive events
                break
            if (neg_event==len(target_list.index) and pos_event==0):
                rank+=-10
                break #when all events are found in negative events

                #if non of above
                #find the closets event to target as the class
                #finally, we do voting to change the rank
            class_offset=np.asarray(class_list, dtype=np.float32)[:,0]
            dist=abs(class_offset-target.offset)
            voting_list.append(class_list[np.argmin(dist)][1])
            if(len(voting_list)==len(target_list.index)): # run when all of the targets in sentence assigned to a class -> voting
                if(sum(voting_list)>0):
                    rank+=10 #voting to positive
                    break
                elif(sum(voting_list)<0):
                    rank+=-10
                    break
                else:
                    rank+=0
                    break

        self.rank+=rank

    def multiplication_score(self, depgraphs, single_id):
        self.tfrank=fn.dep_score_source(depgraphs,single_id,self.Positive,self.Negative)
        self.tgrank=fn.dep_score_target(depgraphs,single_id,self.Positive,self.Negative)

        if(self.tgrank==0 and self.tfrank==0):
            self.rank+=0
            return
        if(self.tgrank==0 and self.tfrank!=0):
            if(self.tfrank>0):
                self.rank+=10
            else:
                self.rank+=-10
            return
        if(self.tgrank!=0 and self.tfrank==0):
            if(self.tgrank>0):
                self.rank+=10
            else:
                self.rank+=-10
            return

        if(self.tgrank!=0 and self.tfrank!=0):
            if(self.tgrank>0 and self.tfrank>0):
                self.rank+=10
                return

            if(self.tgrank<0 and self.tfrank<0):
                self.rank+=10
                return

            if(self.tgrank>0 and self.tfrank<0):
                self.rank+=-10
                return

            if(self.tgrank<0 and self.tfrank>0):
                self.rank+=-10
                return


from warnings import filterwarnings
filterwarnings('ignore')

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import tensorflow as tf
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

from mtcnn import MTCNN
import numpy as np
from PIL import Image
import cv2

class FaceVerification:
    """
    A class for facial recognition with Google's Facenet model embeddings
    """

    def __init__(self, face_detector, model):
        super().__init__()
        self.detector = face_detector
        self.model = model

    def _adjust_gamma(self, image, gamma=2):
        # build a lookup table mapping the pixel values [0, 255] to
        # their adjusted gamma values
        invGamma = 1.0 / gamma
        table = np.array([((i / 255.0)**invGamma) * 255
                          for i in np.arange(0, 256)]).astype("uint8")
        # apply gamma correction using the lookup table
        return cv2.LUT(image, table)

    def extract_face(self,
                     filename,
                     required_size=(160, 160),
                     show_extracted_face=False):
        """
        Extract a single face from a given photograph
        """
        # load image from file
        resized = cv2.cvtColor(filename, cv2.COLOR_BGR2RGB)
        # print(f'Dim Before resizing: {pixels.shape}')

        # Compressing image
        scale_percent = 50  # percent of original size
        width = int(resized.shape[1] * scale_percent / 100)
        height = int(resized.shape[0] * scale_percent / 100)
        dim = (width, height)

        resized = cv2.resize(resized, dim, interpolation=cv2.INTER_AREA)

        # detect faces in the image
        results = self.detector.detect_faces(resized)

        if len(results) > 1:  # image contains more than 1 faces
            sizes = dict()
            for i, res in enumerate(results):
                x1, y1, width, height = res['box']
                area = width * height
                sizes[i] = area
            max_key = [max(sizes, key=sizes.get)][0]
            # Biggest bounding box
            x1, y1, width, height = results[max_key]['box']
            x1, y1 = abs(x1), abs(y1)
            x2, y2 = x1 + width, y1 + height
            face = resized[y1:y2, x1:x2]

        elif len(results) == 1:  # image contains exactly 1 face
            x1, y1, width, height = results[0]['box']
            x1, y1 = abs(x1), abs(y1)
            x2, y2 = x1 + width, y1 + height
            face = resized[y1:y2, x1:x2]

        else:  # Face not detected because of backlight
            resized = cv2.convertScaleAbs(resized, alpha=1.1, beta=2)
            resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
            resized = resized[30::]

            # detect faces in the image
            results = self.detector.detect_faces(resized)

            if len(results) < 1:
                face = np.random.rand(160, 160, 3) * 255
            else:
                # extract the bounding box from the first face
                x1, y1, width, height = results[0]['box']
                x2, y2 = x1 + width, y1 + height

                # extract the face
                face = resized[y1:y2, x1:x2]

        # resize pixels to the model size
        face_array = cv2.resize(face,
                                required_size,
                                interpolation=cv2.INTER_LINEAR_EXACT)
        if show_extracted_face:
            cv2.imshow('face', face_array)
            cv2.waitKey(0)
            cv2.destroyAllWindows()

        face_array = np.expand_dims(face_array, axis=0)
        face_array = face_array / 255
        return face_array

    def compare_distance(self,
                         embedding1=0.0,
                         embedding2=0.0,
                         thresholds=None,
                         distance_metric=None):

        cosine_threshold, euclidean_threshold, euclidean_l2_threshold = thresholds
        # calculate distance with given metric
        if distance_metric == 'euclidean':
            euclidean_distance = round(
                np.linalg.norm(embedding1 - embedding2), 4)
            if euclidean_distance <= euclidean_threshold:
                return 1
            else:
                return 0

        elif distance_metric == 'cosine':
            cosine_distance = (embedding1 @ embedding2.T) / (
                np.linalg.norm(embedding1) * np.linalg.norm(embedding2))
            if cosine_distance <= cosine_threshold:
                return 1
            else:
                return 0

        elif distance_metric == 'euclidean_l2':
            norm1 = embedding1 / np.sqrt(
                np.sum(np.multiply(embedding1, embedding1)))
            norm2 = embedding2 / np.sqrt(
                np.sum(np.multiply(embedding2, embedding2)))
            euclidean_l2_distance = np.linalg.norm(norm1 - norm2)
            if euclidean_l2_distance <= euclidean_l2_threshold:
                return 1
            else:
                return 0

    def extract_and_compare(self,
                            file1=None,
                            file2=None,
                            model=None,
                            thresholds=(0.4, 11.5, 0.8),
                            distance_metric='euclidean'):
        # extract faces
        img1_face = self.extract_face(filename=file1)
        img2_face = self.extract_face(filename=file2)

        # Detect embeddings
        embd1 = self.model.predict(img1_face)
        embd2 = self.model.predict(img2_face)

        # Compute distances and verify
        result = self.compare_distance(embedding1=embd1,
                                       embedding2=embd2,
                                       thresholds=thresholds,
                                       distance_metric=distance_metric)
        return result

if __name__ == '__main__':
    import argparse
    arg_parser = argparse.ArgumentParser(
        'Script for facial recognition threshold testing\n')
    arg_parser.add_argument('--model_path',
                            help='Str, Path to the Model for calculating facial embeddings',
                            required=True)
    arg_parser.add_argument(
        '--image_1', help='Str, Path of the first image', required=True)
    arg_parser.add_argument(
        '--image_2', help='Str, Path of the second image', required=True)
    arg_parser.add_argument(
        '--distance_metric', help='Str, Distance metric to use for verification of images, '
        'could be of the following, "cosine", "euclidean"(default), "euclidean_l2" ')

    args = arg_parser.parse_args()

    # Load face detector
    detector = MTCNN()

    # Load model
    tf.keras.backend.clear_session()
    model = tf.keras.models.load_model(args.model_path, compile=False)

    face_rec = FaceVerification(face_detector=detector, model=model)
    image_1 = cv2.imread(args.image_1)
    image_2 = cv2.imread(args.image_2)
    result = face_rec.extract_and_compare(
        file1=image_1, file2=image_2)
    print('Matched' if result == 1 else "Not matched")


#import required libraries
import keras
from keras.models import load_model
import cv2
import numpy as np
import operator
model=load_model("path/to/your/directory/which/has/model-file/my_model1.h5")

#initialise the camera
cap=cv2.VideoCapture(0)

#declaring the categorical output
categories={0:"call_me",1:"finger_cross",2:"okay",3:"paper",4:"peace",5:"rock",6:"rock_on",7:"scissor",8:"thumbs",9:"up"}

#initialise infinite loop
while True:
       _,frame=cap.read() #reading video from a camera
       frame=cv2.flip(frame,1)
       x1=int(0.5*frame.shape[1]) #to create a ROI on the screen
       y1=10
       x2=frame.shape[1]-10
       y2=int(0.5*frame.shape[1])

       cv2.rectangle(frame,(x1-1,y1-1),(x2-1,y2-1),(255,0,0),1)
       roi=frame[y1:y2,x1:x2]
       roi=cv2.resize(roi,(240,200)) #resize the image
       roi=cv2.cvtColor(roi,cv2.COLOR_BGR2GRAY) #converting BGR to gray scale image
       _,test_image=cv2.threshold(roi,120,255,cv2.THRESH_BINARY_INV) #binary inverse thresholded image
       cv2.imshow("test",test_image)
       result=model.predict(test_image.reshape(1,200,240,1)) #making prediction from the loaded model

       prediction={
                    "call_me":result[0][0],
                    "finger_cross":result[0][1],
                    "okay":result[0][2],
                    "paper":result[0][3],
                    "peace":result[0][4],
                    "rock":result[0][5],
                    "rock_on":result[0][6],
                    "scissor":result[0][7],
                    "thumbs":result[0][8],
                    "up":result[0][9],
           }
       prediction=sorted(prediction.items(),key=operator.itemgetter(1),reverse=True) #sorting of prediction on the basis of max accuracy
       cv2.putText(frame,prediction[0][0],(x1+100,y2+30),cv2.FONT_HERSHEY_PLAIN,2,(0,255,0),3) #showning text on the screen
       cv2.imshow("Frame",frame)

       key=cv2.waitKey(10)
       if key & 0xFF == 27: #press esc for break
           break

cap.release() #switch off camera
cv2.destroyAllWindows() #destroy camera windows


#!/usr/bin/env python
"""
Segment subject for selected thalamic nuclei using whole-brain registration via a template and PICSL label fusion.
"""
import os
import sys
import argparse
import tempfile
import time
import libraries.parallel as parallel
from shutil import rmtree
from functools import partial
from datetime import timedelta
from libraries.imgtools_v0 import check_run, check_warps, sanitize_input, flip_lr, label_fusion_picsl_ants, label_fusion_picsl, ants_compose_a_to_b, ants_apply_only_warp, ants_WarpImageMultiTransform, crop_by_mask, label_fusion_majority
from libraries.ants_nonlinear_v0 import ants_nonlinear_registration, bias_correct, ants_linear_registration
from THOMAS_constants_v0 import image_name, orig_template, template_93, mask_93, this_path, prior_path, subjects, roi, roi_choices, optimal
import nibabel
import numpy as np

def warp_atlas_subject(subject, path, labels, input_image, input_transform_prefix, output_path, exec_options={}):
    """
    Warp a training set subject's labels to input_image.
    """
    a_transform_prefix = os.path.join(path, subject + '/WMnMPRAGE')
    output_path = os.path.join(output_path, subject)
    try:
        os.mkdir(output_path)
    except OSError:
        # Exists
        pass
    combined_warp = os.path.join(output_path, 'Warp.nii.gz')
    if not os.path.exists(combined_warp):
        check_run(
            combined_warp,
            ants_compose_a_to_b,
            a_transform_prefix,
            b_path=input_image,
            b_transform_prefix=input_transform_prefix,
            output=combined_warp,
            **exec_options
        )
    output_labels = {}
    # OPT parallelize, or merge parallelism with subject level
    for label in labels:
        label_fname = os.path.join(path, subject, 'sanitized_rois', label + '.nii.gz')
        warped_label = os.path.join(output_path, label + '.nii.gz')
        switches = '--use-NN'
        check_run(
            warped_label,
            ants_apply_only_warp,
            template=input_image,
            input_image=label_fname,
            input_warp=combined_warp,
            output_image=warped_label,
            switches=switches,
            **exec_options
        )
        output_labels[label] = warped_label
    # Warp anatomical WMnMPRAGE_bias_corr too
    # TODO merge this into previous for loop to be DRY?
    output_labels['WMnMPRAGE_bias_corr'] = output_image = os.path.join(output_path, image_name)
    if not os.path.exists(output_labels['WMnMPRAGE_bias_corr']):
        print output_labels['WMnMPRAGE_bias_corr']
        check_run(
            output_image,
            ants_apply_only_warp,
            template=input_image,
            input_image=os.path.join(path, subject, image_name),
            input_warp=combined_warp,
            output_image=output_image,
            switches='--use-BSpline',
            **exec_options
        )
    return output_labels

def conservative_mask(input_masks, output_path, dilation=0, fill=False):
    """
    Estimates a conservative maximum mask given a list of input masks.
    - for dilation > 0 and fill=True, each side is padded by dilation instead
    - fill will fill the bounding box of the mask producing a cube
    """
    # Maximum label fusion
    # Taken from cv_registration_method.ants_label_fusions
    # cmd = 'AverageImages 3 %s 0 %s' % (output_path, ' '.join(input_masks))
    # cmd = 'ThresholdImage 3 %s %s 0.01 1000' % (output_path, output_path)
    cmd = 'c3d %s -accum -max -endaccum -binarize -o %s' % (' '.join(input_masks), output_path)
    # cmd = 'fslmaths %s -bin %s' % (' -add '.join(input_masks), output_path)
    if sys.platform == 'linux2' or sys.platform == 'darwin':
        parallel_command(cmd)
    if fill:
        # get bounding box
        bbox = map(int, os.popen('fslstats %s -w' % output_path).read().strip().split())
        padding = (-dilation, 2 * dilation) * 3  # min index, size change for 3 spatial dimensions
        if dilation > 0:
            # edit bounding box
            for i, inc in enumerate(padding):  # ignore time dimensions
                bbox[i] += inc
        roi = ' '.join(map(str, bbox))
        # fill bounding box
        cmd = 'fslmaths %s -add 1 -bin -roi %s %s' % (output_path, roi, output_path)
    elif dilation > 0:
        kernel = '%dx%dx%dvox' % (dilation, dilation, dilation)
        cmd = 'c3d %s -dilate 1 %s -o %s' % (output_path, kernel, output_path)
    if sys.platform == 'linux2' or sys.platform == 'darwin':
        parallel_command(cmd)
    return output_path

def get_bounding_box(A):
    B = np.argwhere(A)
    start, stop = B.min(0), B.max(0) + 1
    return zip(start, stop)

def split_roi(roi, axis, split_axis):
    """
    Pages through the roi along axis and cuts each slice in half in the split_axis dimension.
    axis=None cuts the 3D bounding box in half.
    """
    first = np.zeros_like(roi)
    second = np.zeros_like(roi)
    def split_halves(roi, first, second, sl, axis):
        N = len(roi.shape)
        idx = [slice(sl, sl+1) if el is axis else slice(None) for el in xrange(N)]
        try:
            box = get_bounding_box(roi[idx])
        except ValueError:
            # No ROI in this slice
            return
        try:
            box[axis] = tuple(el+sl for el in box[axis])
        except TypeError:
            # Occurs for axis=None case
            pass
        first_idx = [slice(a, a + (b-a)/2) if i is split_axis else slice(a, b) for i, (a, b) in enumerate(box)]
        second_idx = [slice(a + (b-a)/2, b) if i is split_axis else slice(a, b) for i, (a, b) in enumerate(box)]
        # Try pasting in the original ROI to the half boxes
        try:
            first[first_idx] = roi[first_idx]
        except ValueError:
            # exception if half box is 0 along one dimension
            pass
        try:
            second[second_idx] = roi[second_idx]
        except ValueError:
            pass
    if axis is None:
        split_halves(roi, first, second, 0, axis)
    else:
        for sl in xrange(roi.shape[axis]):
            split_halves(roi, first, second, sl, axis)
    return first, second

parser = argparse.ArgumentParser(description='Shortened Template and THalamus for Optimal Multi Atlas Segmentation (ST THOMAS) for a given WMnMPRAGE image using Majority Voting or Joint Label Fusion and the Tourdias atlas.')
parser.add_argument('input_image', help='input WMnMPRAGE NiFTI image, may need to be in LR PA IS format')
#parser.add_argument('output_path', help='the output file for single ROI or directory for multiple ROIs')
parser.add_argument('roi_names', metavar='roi_names', choices=roi_choices, nargs='+', help='a space separated list of one or more ROIs.  Valid targets are: %s' % ', '.join(roi_choices))
parser.add_argument('-a', '--algorithm', type=str, required=True, help='version of THOMAS: v0 or v1 or v2')
parser.add_argument('-w', '--warp', metavar='path', help='looks for {path}InverseWarp.nii.gz and {path}Affine.txt instead of basing it off input_image.')
parser.add_argument('-F', '--forcereg', action='store_true', help='force ANTS registration to WMnMPRAGE mean brain template. The --warp argument can be then used to specify the output path.')
parser.add_argument('-p', '--processes', nargs='?', default=None, const=None, type=int, help='number of parallel processes to use.  If unspecified, automatically set to number of CPUs.')
parser.add_argument('-v', '--verbose', action='store_true', help='verbose mode')
parser.add_argument('-d', '--debug', action='store_true', help='debug mode, interactive prompts')
parser.add_argument('-R', '--right', action='store_true', help='segment right thalamus')
parser.add_argument('-M', '--majorityvoting', action='store_true', help='use majority voting for joint fusion')
parser.add_argument('--jointfusion', action='store_true', help='use older jointfusion instead of antsJointFusion')
parser.add_argument('--tempdir', help='temporary directory to store registered atlases.  This will not be deleted as usual.')
parser.add_argument('--mask', help='custom mask if 93x187x68 mask size is not wanted')
parser.add_argument('--template', help='custom template if 93x187x68 size is not wanted')
parser.add_argument('--output_path', help='specify a different output_path for the output file for single ROI or directory for multiple ROIs')
# TODO handle single roi, single output file case
# TODO fix verbose and debug
# TODO go back to shell=False for command to suppress output and then fix sanitize labels

def main(args, temp_path, pool):
    input_image = orig_input_image = args.input_image

    # assigning default value of mask
    mask = mask_93

    #setting up output path
    if args.output_path:
        output_path = args.output_path
    else:
        output_path = os.path.dirname(orig_input_image)

    #setting up the ROIs
    if roi['param_all'] in args.roi_names:
        labels = list(roi['label_names'])
    else:
        roi_dict = dict(zip(roi['param_names'], roi['label_names']))
        labels = [roi_dict[el] for el in args.roi_names]

    #setting up the template
    if args.algorithm == "v2":
        if args.template is not None and args.mask is not None:
            template = args.template
            mask = args.mask
            print "Custom template and mask"
        elif args.template is not None and args.mask is None:
            sys.exit("!!!!!!! Both template and mask need to be specified simultaneously and they need to be of the same size !!!!!!!")
        elif args.template is None and args.mask is not None:
            sys.exit("!!!!!!! Both template and mask need to be specified simultaneously and they need to be of the same size !!!!!!!")
        else:
            template = template_93
            mask = mask_93
            print "Algorithm is v2"
    elif args.algorithm == "v1":
        sys.exit("!!!!!!! v1 algorithm not yet implemented !!!!!!!")
    elif args.algorithm == "v0":
        template = orig_template
        print "Template is origtemplate.nii.gz"
    else:
        sys.exit("!!!!!!! Algorithm incorrectly specified !!!!!!!")

    # print 'Template being used is'
    # print os.path.abspath(template)

    # TODO prevent both jointfusion and majority voting being set
	# if args.jointfusion is None:
		# print "args.jointfusion has been set (value is %s)" % args.jointfusion
		# if args.majorityvoting is None:
			# print "args.majorityvoting has been set (value is %s)" % args.majorityvoting
			# sys.exit("!!!!!!! Only one label fusion can be selected at any time (default is antsJointFusion) !!!!!!!")

    if args.warp:
        warp_path = args.warp
    else:
        # TODO remove this as the default behavior, instead do ANTS?
        head, tail = os.path.split(input_image)
        tail = tail.replace('.nii', '').replace('.gz', '') #split('.', 1)[0]
        warp_path = os.path.join(temp_path, tail)

    t = time.time()

    if args.algorithm == "v2":
        # Crop the input
        # Affine registering input to template
        print "1.   Linear Registration of input and template \n"
        t1 = time.time()
        ants_linear_registration(orig_template, orig_input_image)
        print '1.   ----- Time Elapsed: %s  \n \n' % timedelta(seconds=time.time()-t1)

        mask_input = os.path.join(os.path.dirname(orig_input_image), 'mask_inp.nii.gz')

        print "2.   Transform mask from template space to input space \n"
        t1 = time.time()
        # Transform mask from template space to input space
        ants_WarpImageMultiTransform(mask, mask_input, orig_input_image)
        print '2.   ----- Time Elapsed: %s  \n \n' % timedelta(seconds=time.time()-t1)

        file_name = os.path.basename(orig_input_image)
        index_of_dot = file_name.index('.')
        file_name_without_extension = file_name[:index_of_dot]
        input_image = os.path.join(os.path.dirname(orig_input_image), 'crop_'+file_name_without_extension+'.nii.gz')

        # Cropping input using this mask
        print "3.   Cropping input using this mask \n"
        t1 = time.time()
        parallel_command(crop_by_mask(orig_input_image, input_image, mask_input))
        print '3.   ----- Time Elapsed: %s  \n \n' % timedelta(seconds=time.time()-t1)

    # FSL automatically converts .nii to .nii.gz
    sanitized_image = os.path.join(temp_path, os.path.basename(input_image) + ('.gz' if input_image.endswith('.nii') else ''))
    if not os.path.exists(sanitized_image):
        input_image = sanitize_input(input_image, sanitized_image, parallel_command)
        if args.right:

            print "4.   lipping along L-R \n"
            t1 = time.time()
            flip_lr(input_image, input_image, parallel_command)
            print '4.   ----- Time Elapsed: %s  \n \n' % timedelta(seconds=time.time()-t1)

        print "5.   Correcting bias \n"
        t1 = time.time()
        bias_correct(input_image, input_image, **exec_options)
        print '5.   ----- Time Elapsed: %s  \n \n' % timedelta(seconds=time.time()-t1)

    else:
        print 'Skipped, using %s' % sanitized_image
        input_image = sanitized_image

    if args.forcereg or not check_warps(warp_path):
        if args.warp:
            print 'Saving output as %s' % warp_path
        else:
            warp_path = os.path.join(temp_path, tail)
            print 'Saving output to temporary path.'
        ants_nonlinear_registration(template, input_image, warp_path, **exec_options)
    else:
        print 'Skipped, using %sInverseWarp.nii.gz and %sAffine.txt' % (warp_path, warp_path)

    # generating the warped output
    print "6.   Warping prior labels and images \n \n"
    t1 = time.time()
    registered = os.path.join(temp_path, 'registered.nii.gz')
    cmd = 'WarpImageMultiTransform 3 %s %s -R %s %sWarp.nii.gz %sAffine.txt' % (input_image, registered, template, warp_path, warp_path)
    parallel_command(cmd)

    print '6.   --- Elapsed: %s  \n \n' % timedelta(seconds=time.time()-t1)

    t1 = time.time()
    # TODO should probably use output from warp_atlas_subject instead of hard coding paths in create_atlas
    # TODO make this more parallel
    warped_labels = pool.map(partial(
        warp_atlas_subject,
        path=prior_path,
        # TODO cleanup this hack to always have whole thalamus so can estimate mask
        labels=set(labels + ['1-THALAMUS']),
        input_image=input_image,
        input_transform_prefix=warp_path,
        output_path=temp_path,
        exec_options=exec_options,
    ), subjects)
    warped_labels = {label: {subj: d[label] for subj, d in zip(subjects, warped_labels)} for label in warped_labels[0]}
    # # print '--- Forming subject-registered atlases. --- Elapsed: %s' % timedelta(seconds=time.time()-t)
    # atlases = pool.map(partial(create_atlas, path=temp_path, subjects=subjects, target='', echo=exec_options['echo']),
    # [{'label': label, 'output_atlas': os.path.join(temp_path, label+'_atlas.nii.gz')} for label in warped_labels])
    # atlases = dict(zip(warped_labels, zip(*atlases)[0]))
    # atlas_image = atlases['WMnMPRAGE_bias_corr']
    atlas_images = warped_labels['WMnMPRAGE_bias_corr'].values()

    print '!! --- Performing label fusion. --- Elapsed: %s  \n \n' % timedelta(seconds=time.time() - t1)
    # FIXME use whole-brain template registration optimized parameters instead, these are from crop pipeline
    optimal_picsl = optimal['PICSL']
    # for k, v in warped_labels.iteritems():
    #     print k, v
    # for label in labels:
    #     print optimal_picsl[label]
    if args.jointfusion:
        pool.map(partial(label_fusion_picsl, input_image, atlas_images),
                 [dict(
                     atlas_labels=warped_labels[label].values(),
                     output_label=os.path.join(temp_path, label + '.nii.gz'),
                     rp=optimal_picsl[label]['rp'],
                     rs=optimal_picsl[label]['rs'],
                     beta=optimal_picsl[label]['beta'],
                     **exec_options
                 ) for label in labels])
    elif args.majorityvoting:
        pool.map(partial(label_fusion_majority),
                 [dict(
                     atlas_labels=warped_labels[label].values(),
                     output_label=os.path.join(temp_path, label + '.nii.gz'),
                     rp=optimal_picsl[label]['rp'],
                     rs=optimal_picsl[label]['rs'],
                     beta=optimal_picsl[label]['beta'],
                     **exec_options
                 ) for label in labels])
    else:
        # Estimate mask to restrict computation
        mask = os.path.join(temp_path, 'mask.nii.gz')
        check_run(
            mask,
            conservative_mask,
            warped_labels['1-THALAMUS'].values(),
            mask,
            dilation=10,
        )
        pool.map(partial(label_fusion_picsl_ants, input_image, atlas_images),
                 [dict(
                     atlas_labels=warped_labels[label].values(),
                     output_label=os.path.join(temp_path, label + '.nii.gz'),
                     rp=optimal_picsl[label]['rp'],
                     rs=optimal_picsl[label]['rs'],
                     beta=optimal_picsl[label]['beta'],
                     mask=mask,
                     **exec_options
                 ) for label in labels])

    print '!! --- Performing label fusion. --- Elapsed: %s  \n \n' % timedelta(seconds=time.time() - t1)

    # STEPS
    # pool_small.map(partial(label_fusion, input_image=input_image, image_atlas=atlases['WMnMPRAGE_bias_corr'], echo=exec_options['echo']),
    #     [{
    #         'label_atlas': atlases[label],
    #         'output_label': os.path.join(output_path, label+'.nii.gz'),
    #         'sigma': optimal_steps[label]['steps_sigma'],
    #         'X': optimal_steps[label]['steps_X'],
    #         'mrf': optimal_steps[label]['steps_mrf'],
    #     } for label in labels]
    # )
    # for label in labels:
    #     print {
    #         'label': label,
    #         'sigma': optimal_steps[label]['steps_sigma'],
    #         'X': optimal_steps[label]['steps_X'],
    #         'mrf': optimal_steps[label]['steps_mrf'],
    #     }
    #     partial_fusion = partial(label_fusion, input_image=input_image, image_atlas=atlases['WMnMPRAGE_bias_corr'], echo=exec_options['echo'])
    #     label_fusion_args = {
    #         'label_atlas': atlases[label],
    #         'output_label': os.path.join(output_path, label+'.nii.gz'),
    #         'sigma': optimal_steps[label]['steps_sigma'],
    #         'X': optimal_steps[label]['steps_X'],
    #         'mrf': optimal_steps[label]['steps_mrf'],
    #     }
    #     partial_fusion(**label_fusion_args)

    files = [(os.path.join(temp_path, label + '.nii.gz'), os.path.join(output_path, label + '.nii.gz')) for label in labels]
    if args.right:
        pool.map(flip_lr, files)
        files = [(os.path.join(output_path, label + '.nii.gz'), os.path.join(output_path, label + '.nii.gz')) for label in labels]
    # Resort output to original ordering
    pool.map(parallel_command,
        ['%s %s %s %s' % (os.path.join(this_path, 'swapdimlike.py'), in_file, orig_input_image, out_file) for in_file, out_file in files])

    # get the vlp file path for splitting
    vlp_file = os.path.join(output_path, '6-VLP.nii.gz')

    # Re-orient to standard space - LR PA IS format
    san_vlp_file = os.path.join(output_path, 'san_6-VLP.nii.gz')
    input_image1 = sanitize_input(vlp_file, san_vlp_file, parallel_command)

    # get the sanitized vlp for processing
    input_nii = nibabel.load(input_image1)
    data = input_nii.get_data()
    hdr = input_nii.get_header()
    affine = input_nii.get_affine()

    # Coronal axis for RL PA IS orientation
    vlps = split_roi(data, None, 2)
    for fname, sub_vlp in zip(['6_VLPv.nii.gz', '6_VLPd.nii.gz'], vlps):
        output_nii = nibabel.Nifti1Image(sub_vlp, affine, hdr)
        output_nii.to_filename(os.path.join(os.path.dirname(out_file), fname))

    print '--- Finished --- Elapsed: %s' % timedelta(seconds=time.time() - t)

if __name__ == '__main__':
    args = parser.parse_args()
    # print args
    # exec_options.update({'debug': args.debug, 'verbose': args.verbose})
    exec_options = {'echo': False, 'suppress': True}
    if args.verbose:
        exec_options['verbose'] = True
    if args.debug:
        print 'Debugging mode forces serial execution.'
        # exec_options['echo'] = True
        args.processes = 1
    parallel_command = partial(parallel.command, **exec_options)
    pool = parallel.BetterPool(args.processes)
    print 'Running with %d processes.' % pool._processes
    # TODO don't hard code this number of processors
    # pool_small = parallel.BetterPool(4)
    # TODO Add path of script to command()
    # os.environ['PATH'] += os.pathsep + os.path.abspath(os.path.dirname(sys.argv[0]))
    if args.tempdir:
        temp_path = args.tempdir
        if not os.path.exists(temp_path):
            print 'Making %s' % os.path.abspath(temp_path)
            os.makedirs(temp_path)
    else:
        temp_path = tempfile.mkdtemp(dir=os.path.dirname(args.output_path))
    try:
        main(args, temp_path, pool)
    finally:
        pool.close()
        # Clean up temp folders
        if not args.debug and not args.tempdir:
            try:
                rmtree(temp_path)
            except OSError as exc:
                if exc.errno != 2:  # Code 2 - no such file or directory
                    raise


from ctypes import cdll, c_double
# from ctypes.util import find_library
import ctypes
# from numpy.ctypeslib import as_ctypes, as_array, as_ctypes_type
import numpy as np
from numpy.ctypeslib import ndpointer
import copy
import logging
from .basefunctions import find_lib, get_root_path
import os

# relative path to colorprocesser library
RELATIVE_LIBRARY_PATH = "build/"

# lib_path = find_library('colorprocesser')
lib_path = find_lib(os.path.join(get_root_path(), RELATIVE_LIBRARY_PATH), 'libcolorprocesser')
if not bool(lib_path):
    raise FileNotFoundError("colorprocesser library not found!")
# lib_full_path = os.path.join(get_root_path(), RELATIVE_LIBRARY_PATH, lib_path)
lib_full_path = lib_path
logging.info("colorprocesser library path: {0}".format(lib_full_path))
lib = cdll.LoadLibrary(lib_full_path)
# set output types for ColorProcesser methods
lib.ColorProcesser_helloworld.restype = c_double
lib.ColorProcesser_test_calc.restype = c_double

lib.ColorProcesser_color_distance.restype = c_double

class ColorProcesser(object):
    def __init__(self):
        self.obj = lib.ColorProcesser_new()
    #     fun = lib.ColorProcesser_new
    #     fun.argtypes = []
    #     fun.restype = ctypes.c_void_p
    #     self.obj = fun()
    #
    # def __del__(self):
    #     fun = lib.ColorProcesser_delete
    #     fun.argtypes = [ctypes.c_void_p]
    #     fun.restype = None
    #     fun(self.obj)

    def helloworld(self):
        return lib.ColorProcesser_helloworld(self.obj)

    def test_calc(self):
        return lib.ColorProcesser_test_calc(self.obj)

    def color_distance(self, c1, c2):

        return lib.ColorProcesser_color_distance(
            self.obj,
            ctypes.c_double(c1[0]),
            ctypes.c_double(c1[1]),
            ctypes.c_double(c1[2]),
            ctypes.c_double(c2[0]),
            ctypes.c_double(c2[1]),
            ctypes.c_double(c2[2])
        )

    def array_color_distance(self, the_color, color_array):

        input_shape = color_array.shape

        # print(input_shape)

        m = np.prod(input_shape[:-1])

        # print(m)

        new_color_array = color_array.reshape(m, 3)

        # print(new_color_array)

        r1 = the_color[0]
        g1 = the_color[1]
        b1 = the_color[2]

        r2 = np.ascontiguousarray(new_color_array[:, 0], dtype=np.double)
        g2 = np.ascontiguousarray(new_color_array[:, 1], dtype=np.double)
        b2 = np.ascontiguousarray(new_color_array[:, 2], dtype=np.double)

        # the_color_distances = np.ascontiguousarray(np.zeros((m,)), dtype=np.double)

        lib.ColorProcesser_array_color_distance.restype = ndpointer(dtype=c_double, shape=(m, ))

        res_array = lib.ColorProcesser_array_color_distance(
            self.obj,
            c_double(r1),
            c_double(g1),
            c_double(b1),
            ctypes.c_void_p(r2.ctypes.data),
            ctypes.c_void_p(g2.ctypes.data),
            ctypes.c_void_p(b2.ctypes.data),
            # ctypes.c_void_p(the_color_distances.ctypes.data),
            ctypes.c_int(m)
        )

        # res_array_content = ctypes.cast(res_array, ndpointer(dtype=c_double, shape=(m, )))

        res_array_copy = copy.deepcopy(res_array)

        # res_array_p = ctypes.cast(res_array, ctypes.POINTER(c_double))

        res_array_p = res_array.ctypes.data_as(ctypes.POINTER(c_double))

        lib.free_double_array(res_array_p)

        return res_array_copy.reshape(input_shape[:-1])


import os.path
import importlib
from pathlib import Path

import bcolz

from scipy.interpolate import interp1d
from scipy import signal
import numpy as np
import pandas as pd

import logging
log = logging.getLogger(__name__)

from atom.api import Atom, Callable, Dict, Float, Property, Typed, Value
from enaml.core.api import Declarative, d_
from psi import SimpleState

from . import util

from psi import get_config

################################################################################
# Exceptions
################################################################################
mesg = '''
Unable to run the calibration. Please double-check that the microphone and
speaker amplifiers are powered on and that the devices are positioned properly.
If you keep receiving this message, the microphone and/or the speaker may have
gone bad and need to be replaced.

{}
'''
mesg = mesg.strip()

thd_err_mesg = 'Total harmonic distortion for {:.1f}Hz is {:.1f}%'
nf_err_mesg = 'Power at {:.1f}Hz has SNR of {:.2f}dB'

class CalibrationError(Exception):

    def __str__(self):
        return self.message

class CalibrationTHDError(CalibrationError):

    def __init__(self, frequency, thd):
        self.frequency = frequency
        self.thd = thd
        self.base_message = thd_err_mesg.format(frequency, thd)
        self.message = mesg.format(self.base_message)

class CalibrationNFError(CalibrationError):

    def __init__(self, frequency, snr):
        self.frequency = frequency
        self.snr = snr
        self.base_message = nf_err_mesg.format(frequency, snr)
        self.message = mesg.format(self.base_message)

################################################################################
# Calibration loaders
################################################################################
def from_psi_chirp(folder, output_gain=None, **kwargs):
    filename = os.path.join(folder, 'chirp_sensitivity.csv')
    sensitivity = pd.io.parsers.read_csv(filename)
    if output_gain is None:
        output_gain = sensitivity.loc[:, 'hw_ao_chirp_level'].max()
    m = sensitivity['hw_ao_chirp_level'] == output_gain
    mic_freq = sensitivity.loc[m, 'frequency'].values
    mic_sens = sensitivity.loc[m, 'sens'].values
    return InterpCalibration(mic_freq, mic_sens, **kwargs)

def from_psi_golay(folder, n_bits=None, output_gain=None, **kwargs):
    filename = os.path.join(folder, 'sensitivity.csv')
    sensitivity = pd.io.parsers.read_csv(filename)
    if n_bits is None:
        n_bits = sensitivity['n_bits'].max()
    if output_gain is None:
        m = sensitivity['n_bits'] == n_bits
        output_gain = sensitivity.loc[m, 'output_gain'].max()
    m_n_bits = sensitivity['n_bits'] == n_bits
    m_output_gain = sensitivity['output_gain'] == output_gain
    m = m_n_bits & m_output_gain
    mic_freq = sensitivity.loc[m, 'frequency'].values
    mic_sens = sensitivity.loc[m, 'sens'].values
    return InterpCalibration(mic_freq, mic_sens, **kwargs)

def from_epl(filename, **kwargs):
    calibration = pd.io.parsers.read_csv(filename, skiprows=14, delimiter='\t')
    freq = calibration['Freq(Hz)']
    spl = calibration['Mag(dB)']
    return InterpCalibration.from_spl(freq, spl, **kwargs)

def from_bcolz_store(rootdir):
    '''
    Returns the calibration used by the input
    '''
    import bcolz
    with bcolz.carray(rootdir=rootdir) as fh:
        data = fh.attrs['source']['calibration']
        class_name = data.pop('type')
        cls = globals()[class_name]
        return cls(**data)

def from_tone_sens(folder):
    filename = os.path.join(folder, 'tone_sensitivity.csv')
    df = pd.read_csv(filename)
    return PointCalibration(df['frequency'], df['sens'])

################################################################################
# Calibration routines
################################################################################
class Calibration(Atom):
    '''
    Assumes that the system is linear for a given frequency

    Parameters
    ----------
    frequency : 1D array
        Frequencies that system sensitivity was measured at.
    sensitivity : 1D array
        Sensitivity of system in dB(V/Pa).
    '''
    attrs = Dict().tag(metadata=True)

    @classmethod
    def as_attenuation(cls, vrms=1, **kwargs):
        '''
        Allows levels to be specified in dB attenuation
        '''
        return cls.from_spl([0, 100e3], [0, 0], vrms, **kwargs)

    @classmethod
    def from_magnitude(cls, frequency, magnitude, vrms=1, **kwargs):
        '''
        Generates a calibration object based on the recorded magnitude (Pa)

        Parameters
        ----------
        frequency : array-like
            List of freuquencies (in Hz)
        magnitude : array-like
            List of magnitudes (e.g., speaker output in Pa) for the specified
            RMS voltage.
        vrms : float
            RMS voltage (in Volts)

        Additional kwargs are passed to the class initialization.
        '''
        sensitivity = util.db(vrms)-util.db(magnitude)-util.db(20e-6)
        return cls(frequency, sensitivity, **kwargs)

    @classmethod
    def from_spl(cls, frequency, spl, vrms=1, **kwargs):
        '''
        Generates a calibration object based on the recorded SPL

        Parameters
        ----------
        frequency : array-like
            List of freuquencies (in Hz)
        spl : array-like
            List of magnitudes (e.g., speaker output in SPL) for the specified
            RMS voltage.
        vrms : float
            RMS voltage (in Volts)

        Additional kwargs are passed to the class initialization.
        '''
        sensitivity = util.db(vrms)-spl-util.db(20e-6)
        return cls(frequency, sensitivity, **kwargs)

    def get_spl(self, frequency, voltage):
        sensitivity = self.get_sens(frequency)
        return util.db(voltage)-sensitivity-util.db(20e-6)

    def get_sf(self, frequency, spl, attenuation=0):
        sensitivity = self.get_sens(frequency)
        vdb = sensitivity+spl+util.db(20e-6)+attenuation
        return 10**(vdb/20.0)

    def get_mean_sf(self, flb, fub, spl, attenuation=0):
        frequencies = np.arange(flb, fub)
        return self.get_sf(frequencies, spl).mean(axis=0)

    def get_attenuation(self, frequency, voltage, level):
        return self.get_spl(frequency, voltage)-level

    def get_gain(self, frequency, spl, attenuation=0):
        return util.db(self.get_sf(frequency, spl, attenuation))

    def set_fixed_gain(self, fixed_gain):
        self.fixed_gain = fixed_gain

    def get_sens(self, frequency):
        raise NotImplementedError

class FlatCalibration(Calibration):

    sensitivity = Float().tag(metadata=True)
    fixed_gain = Float(0).tag(metadata=True)
    mv_pa = Property()

    def _get_mv_pa(self):
        return util.dbi(self.sensitivity) * 1e3

    def _set_mv_pa(self, value):
        self.sensitivity = util.db(value * 1e-3)

    @classmethod
    def as_attenuation(cls, vrms=1, **kwargs):
        '''
        Allows levels to be specified in dB attenuation
        '''
        return cls.from_spl(0, vrms, **kwargs)

    @classmethod
    def from_spl(cls, spl, vrms=1, **kwargs):
        '''
        Generates a calibration object based on the recorded SPL

        Parameters
        ----------
        spl : array-like
            List of magnitudes (e.g., speaker output in SPL) for the specified
            RMS voltage.
        vrms : float
            RMS voltage (in Volts)

        Additional kwargs are passed to the class initialization.
        '''
        sensitivity = util.db(vrms)-spl-util.db(20e-6)
        return cls(sensitivity=sensitivity, **kwargs)

    @classmethod
    def from_mv_pa(cls, mv_pa, **kwargs):
        sens = util.db(mv_pa*1e-3)
        return cls(sensitivity=sens, **kwargs)

    def get_sens(self, frequency):
        return self.sensitivity-self.fixed_gain

    def get_mean_sf(self, flb, fub, spl, attenuation=0):
        return self.get_sf(flb, spl)

class UnityCalibration(FlatCalibration):

    # For unity calibration, set the property so it doesn't get saved.  This
    # value gives us unity passthrough (because the core methods assume
    # everything is in units of dB(Vrms/Pa)).
    sensitivity = Float(-util.db(20e-6)).tag(metadata=False)

class InterpCalibration(Calibration):
    '''
    Use when calibration is not flat (i.e., uniform) across frequency.

    Parameters
    ----------
    frequency : array-like, Hz
        Calibrated frequencies (in Hz)
    sensitivity : array-like, dB(V/Pa)
        Sensitivity at calibrated frequency in dB(V/Pa) assuming 1 Vrms and 0 dB
        gain.  If you have sensitivity in V/Pa, just pass it in as
        20*np.log10(sens).
    fixed_gain : float
        Fixed gain of the input or output.  The sensitivity is calculated using
        a gain of 0 dB, so if the input (e.g. a microphone preamp) or output
        (e.g. a speaker amplifier) adds a fixed gain, this needs to be factored
        into the calculation.

        For input calibrations, the gain must be negative (e.g. if the
        microphone amplifier is set to 40 dB gain, then provide -40 as the
        value).
    attrs : {None, dict}
        Extra attrs that will be saved alongside the calibration.
    '''

    frequency = Typed(np.ndarray).tag(metadata=True)
    sensitivity = Typed(np.ndarray).tag(metadata=True)
    fixed_gain = Float(0).tag(metadata=True)
    _interp = Callable()

    def __init__(self, frequency, sensitivity, fixed_gain=0, attrs=None):
        if attrs is None:
            attrs = {}
        self.frequency = np.asarray(frequency)
        self.sensitivity = np.asarray(sensitivity)
        self.fixed_gain = fixed_gain
        self.attrs = attrs
        self._interp = interp1d(frequency, sensitivity, 'linear', bounds_error=False)

    def get_sens(self, frequency):
        # Since sensitivity is in dB(V/Pa), subtracting fixed_gain from
        # sensitivity will *increase* the sensitivity of the system.
        return self._interp(frequency)-self.fixed_gain

class PointCalibration(Calibration):

    frequency = Typed(np.ndarray).tag(metadata=True)
    sensitivity = Typed(np.ndarray).tag(metadata=True)
    fixed_gain = Float(0).tag(metadata=True)

    def __init__(self, frequency, sensitivity, fixed_gain=0, attrs=None):
        if attrs is None:
            attrs = {}
        if np.isscalar(frequency):
            frequency = [frequency]
        if np.isscalar(sensitivity):
            sensitivity = [sensitivity]
        self.frequency = np.array(frequency)
        self.sensitivity = np.array(sensitivity)
        self.fixed_gain = fixed_gain
        self.attrs = attrs

    def get_sens(self, frequency):
        if np.iterable(frequency):
            return np.array([self._get_sens(f) for f in frequency])
        else:
            return self._get_sens(frequency)

    def _get_sens(self, frequency):
        try:
            i = np.flatnonzero(np.equal(self.frequency, frequency))[0]
        except IndexError:
            log.debug('Calibrated frequencies are %r', self.frequency)
            m = 'Frequency {} not calibrated'.format(frequency)
            raise CalibrationError(m)
        return self.sensitivity[i]-self.fixed_gain

    @classmethod
    def from_psi_chirp(cls, folder, output_gain=None, **kwargs):
        filename = os.path.join(folder, 'chirp_sensitivity.csv')
        sensitivity = pd.io.parsers.read_csv(filename)
        if output_gain is None:
            output_gain = sensitivity.loc[:, 'hw_ao_chirp_level'].max()
        m = sensitivity['hw_ao_chirp_level'] == output_gain
        mic_freq = sensitivity.loc[m, 'frequency'].values
        mic_sens = sensitivity.loc[m, 'sens'].values
        attrs = {
            'type': 'psi_chirp',
            'filename': str(folder),
            'output_gain': output_gain,
        }
        return cls(mic_freq, mic_sens, attrs=attrs, **kwargs)

class EPLCalibration(InterpCalibration):

    @classmethod
    def load_data(cls, filename):
        filename = Path(filename)
        calibration = pd.io.parsers.read_csv(filename, skiprows=14,
                                             delimiter='\t')
        freq = calibration['Freq(Hz)']
        spl = calibration['Mag(dB)']
        return {
            'attrs': {
                'source': filename
            },
            'frequency': freq,
            'spl': spl,
        }

    @classmethod
    def load(cls, filename, **kwargs):
        data = cls.load_data(filename)
        data.update(kwargs)
        return cls.from_spl(**data)

class CochlearCalibration(InterpCalibration):

    @classmethod
    def load_data(cls, filename):
        import tables
        with tables.open_file(filename, 'r') as fh:
            mic_freq = np.asarray(fh.get_node('/frequency').read())
            mic_sens = np.asarray(fh.get_node('/exp_mic_sens').read())
            return {
                'attrs': {
                    'source': filename
                },
                'frequency': mic_freq,
                'sensitivity': mic_sens,
            }

    @classmethod
    def load(cls, filename, **kwargs):
        data = cls.load_data(filename)
        data.update(kwargs)
        return cls(**kwargs)

class GolayCalibration(InterpCalibration):

    fs = Float().tag(metadata=True)
    phase = Typed(np.ndarray).tag(metadata=True)

    def __init__(self, frequency, sensitivity, fs=None, phase=None,
                 fixed_gain=0, **kwargs):
        super().__init__(frequency, sensitivity, fixed_gain, **kwargs)
        # fs and phase are required for the IIR stuff
        if fs is not None:
            self.fs = fs
        if phase is not None:
            self.phase = np.asarray(phase)

    @staticmethod
    def load_data(folder, n_bits=None, output_gain=None):
        from psi.data.io.calibration import CalibrationFile
        fh = CalibrationFile(folder)
        return fh._get_golay_data(n_bits, output_gain)

    @classmethod
    def load(cls, folder, n_bits=None, output_gain=None, **kwargs):
        from psi.data.io.calibration import CalibrationFile
        fh = CalibrationFile(folder)
        return fh.get_golay_calibration(n_bits, output_gain)

    def get_iir(self, fs, fl, fh, truncate=None):
        fs_ratio = self.fs/fs
        if int(fs_ratio) != fs_ratio:
            m = 'Calibration sampling rate, {}, must be an ' \
                'integer multiple of the requested sampling rate'
            raise ValueError(m.format(self.fs))

        n = (len(self.frequency)-1)/fs_ratio + 1
        if int(n) != n:
            m = 'Cannot achieve requested sampling rate ' \
                'TODO: explain why'
            raise ValueError(m)
        n = int(n)

        fc = (fl+fh)/2.0
        freq = self.frequency
        phase = self.phase
        sens = self.sensitivity - (self.get_sens(fc) + self.fixed_gain)
        sens[freq < fl] = 0
        sens[freq >= fh] = 0
        m, b = np.polyfit(freq[freq < fh], phase[freq < fh], 1)
        invphase = 2*np.pi*np.arange(len(freq))*m
        inv_csd = util.dbi(sens)*np.exp(invphase*1j)

        # Need to trim so that the data is resampled accordingly
        if fs_ratio != 1:
            inv_csd = inv_csd[:n]
        iir = np.fft.irfft(inv_csd)

        if truncate is not None:
            n = int(truncate*fs)
            iir = iir[:n]

        return iir

class ChirpCalibration(InterpCalibration):

    @staticmethod
    def load_data(folder, output_gain=None):
        folder = Path(folder)
        sensitivity = pd.io.parsers.read_csv(folder / 'chirp_summary.csv')
        if output_gain is None:
            output_gain = sensitivity['hw_ao_chirp_level'].max()

        m = sensitivity['hw_ao_chirp_level'] == output_gain
        mic_freq = sensitivity.loc[m, 'frequency'].values
        mic_sens = sensitivity.loc[m, 'sens'].values
        mic_phase = sensitivity.loc[m, 'phase'].values
        return {
            'attrs': {
                'source': folder,
                'output_gain': output_gain,
                'calibration_type': 'psi_chirp',
            },
            'frequency': mic_freq,
            'sensitivity': mic_sens,
        }

    @classmethod
    def load(cls, folder, n_bits=None, output_gain=None, **kwargs):
        data = cls.load_psi_golay(folder, n_bits, output_gain)
        data.update(kwargs)
        return cls(**data)

class CalibrationRegistry:

    def __init__(self):
        self.registry = {}

    def register(self, klass, label=None):
        calibration_type = klass.__name__
        calibration_path = f'{klass.__module__}.{calibration_type}'
        if label is None:
            label = calibration_type
        if calibration_type in self.registry:
            m = f'{label} already registered as {calibration_type}'
            raise ValueError(m)
        self.registry[calibration_path] = klass, label
        log.debug('Registered %s', calibration_path)

    def clear(self):
        self.registry.clear()

    def register_basic(self, clear=False, unity=True, fixed=True, golay=True,
                       chirp=True):
        if clear:
            self.clear()
        if unity:
            self.register(UnityCalibration, 'unity gain')
        if fixed:
            self.register(FlatCalibration, 'fixed sensitivity')
        if golay:
            self.register(GolayCalibration, 'Golay calibration')
        if chirp:
            self.register(ChirpCalibration, 'Chirp calibration')

    def get_classes(self):
        return [v[0] for v in self.registry.values()]

    def get_class(self, calibration_type):
        return self.registry[calibration_type][0]

    def get_labels(self):
        return [v[1] for v in self.registry.values()]

    def get_label(self, obj):
        name = f'{obj.__module__}.{obj.__name__}'
        return self.registry[name][1]

    def from_dict(self, calibration_type, **kw):
        if calibration_type not in self.registry:
            log.debug('Importing and registering calibration')
            # Older calibration formats may still have only the class name, not
            # the full module + class name.
            try:
                module_name, class_name = calibration_type.rsplit('.', 1)
            except ValueError:
                module_name = __name__
                class_name = calibration_type
            module = importlib.import_module(module_name)
            klass = getattr(module, class_name)
        else:
            klass = self.get_class(calibration_type)
        return klass(**kw)

calibration_registry = CalibrationRegistry()
calibration_registry.register_basic()
calibration_registry.register(EPLCalibration, 'EPL calibration')
calibration_registry.register(CochlearCalibration, 'Golay calibration (old Cochlear format)')

if __name__ == '__main__':
    import doctest
    doctest.testmod()


import numpy as np
import pickle
import cv2
import keras
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import OneHotEncoder, LabelEncoder
from keras.preprocessing.image import ImageDataGenerator

class SF_DataGenerator(keras.utils.Sequence):
    def __init__(self, config, train_flag = False):
        'Initialization'
        self.config = config
        self.batch_size = self.config.batch_size
        self.n_classes = self.config.num_classes
        self.train_flag = train_flag

        if train_flag == True:
            data_pd = pd.read_pickle(self.config.train_csv)
        else:
            data_pd = pd.read_pickle(self.config.valid_csv)
        img_len = self.config.data_dir+'iiiiiiiimgs/train/c0/'+data_pd.iloc[0]['img']
        self.data = np.reshape(np.array([img_len, 'apples']), (1,2))
        self.data = np.repeat(self.data, len(data_pd), axis=0)

        for sf_data_itr in range(len(data_pd)):
            img_name = data_pd.iloc[sf_data_itr]['img']
            c_x = data_pd.iloc[sf_data_itr]['classname']
            self.data[sf_data_itr, :] = (self.config.data_dir+'imgs/train/'+c_x+'/'+img_name,(c_x.split('c')[1]))

        #Iterations and data length
        self.data_len = len(data_pd)
        self.data_iters = ( self.data_len + self.config.batch_size ) // self.config.batch_size

        #Keras Augmtations
        if train_flag == True:
            data_gen_args = dict(rotation_range=0.2,
                width_shift_range=0.05,
                height_shift_range=0.05,
                shear_range=0.05,
                zoom_range=0.05,
                horizontal_flip=True,
                fill_mode='nearest')
            self.augmentor =ImageDataGenerator(**data_gen_args)

    def __len__(self):
        return int(np.ceil(self.data_len / float(self.batch_size)))

    def on_epoch_end(self):
        if self.train_flag == True:
            np.random.shuffle(self.data)

    def __getitem__(self, idx):
        batch_start = idx*self.batch_size
        batch_end = (idx+1)*self.batch_size
        img_batch = []
        ann_batch = []

        for itr in range(batch_start, batch_end):
            if itr >= self.data_len:
                break
            img = cv2.imread(self.data[itr,0])
            img = cv2.resize(img, (self.config.resize[1], self.config.resize[0]))
            ann = keras.utils.to_categorical(int(self.data[itr,1]), num_classes=self.n_classes)
            if self.train_flag == True:
                img = self.augmentor.random_transform(img)
            img_batch.append(img)
            ann_batch.append(ann)
        img_batch = np.array(img_batch)
        ann_batch = np.array(ann_batch)
        return img_batch, ann_batch


import argparse
import sys

import cv2
import tensorflow as tf
import numpy as np
from backend.utils import decode_regression_n

def softmax(x):
    x_row_max = x.max(axis=-1)
    x_row_max = x_row_max.reshape(list(x.shape)[:-1]+[1])
    x = x - x_row_max
    x_exp = np.exp(x)
    x_exp_row_sum = x_exp.sum(axis=-1).reshape(list(x.shape)[:-1]+[1])
    softmax = x_exp / x_exp_row_sum
    return softmax

conf_threshold = 0.6
center_variance = 0.1
size_variance = 0.2

image_size = [320, 240]  # default input size 320*240
feature_map_wh_list = [[10, 8]]  # default feature map size
min_boxes = [[64, 96]]

def main():

    model = tf.keras.models.load_model("models/slim/face.h5")

    img = cv2.imread("imgs/t1.jpg")
    h, w, _ = img.shape
    img_resize = cv2.resize(img, (320, 240))
    img_resize = cv2.cvtColor(img_resize, cv2.COLOR_BGR2RGB)
    img_resize = img_resize - 127.0
    img_resize = img_resize / 128.0

    results = model.predict(np.expand_dims(img_resize, axis=0))  # result=[background,face,x1,y1,x2,y2]
    print(results.shape)
    reg,cls = results[:,:,:,0:8],results[:,:,:,8:]
    reg = reg.reshape([-1,4])
    print(reg.shape)
    cls = cls.reshape([-1, 2])
    print(cls.shape)

   # cls = np.Softmax(axis=-1)(cls)
    cls = softmax(cls)
    print(cls[:,1]>conf_threshold)
    cond = np.where(cls[:,1]>conf_threshold )
    loc = decode_regression_n(reg, image_size, feature_map_wh_list, min_boxes,
                            center_variance, size_variance)
    print(loc[cond])
    results = loc[cond]
    for result in results:
        start_x = int(result[0] * w)
        start_y = int(result[1] * h)
        end_x = int(result[2] * w)
        end_y = int(result[3] * h)

        cv2.rectangle(img, (start_x, start_y), (end_x, end_y), (0, 255, 0), 2)

    cv2.imwrite('imgs/test_output_res.jpg', img)

if __name__ == '__main__':
    main()