#!/usr/bin/python
#    By Hernan Chavez Thielemann
__author__ = 'Hernan Chavez Thielemann <hchavezthiele at gmail dot com>'
# checked ok 30/04/2018

#------------------------------------------------------
#///    Packages and globals definitions are here   ///
#------------------------------------------------------
from os.path import dirname, realpath
from sys import exit

from Tkinter import Tk, Frame, Label, TclError, PhotoImage

from conversion_gui import Conversion
from script_gui import Script_GUI
from run_gui import Run_GUI

from popup import AboutPopUp
from tk_lib import createmenubar

from lib.misc.warn import wrg_3
from lib.misc.file import run_command
from lib.misc.version import __version__

#------------------------------------------------------
'''///////////////        Class        /////////////'''
#------------------------------------------------------

class Gro2Lam_GUI(Frame):
    ''' Graphic User Interface '''
    def __init__(self, master=None, test = False):
        Frame.__init__(self, master)
        _ver= __version__.split()
        self.master.title(" "*5+"{}    {}".format(_ver[0],_ver[2]))#.master

        self.pack() # ... why I'm packing here?? coords?
        self.test = test

        # images storaging
        dir_path = dirname( realpath( __file__))
        self.img = dict()
        self.img['logo']    = PhotoImage( file = dir_path + "/img/logo.ppm")
        self.img['help']    = PhotoImage( file = dir_path + "/img/help.ppm")
        self.img['file']    = PhotoImage( file = dir_path + "/img/file.ppm")
        self.img['gear']    = PhotoImage( file = dir_path + "/img/gear.ppm")

        # body init
        self.prevailing_body = 0
        self.body = None
        self.MAINVERTEX = [ 0, 0, 0, 0, 0, 0]

        # Conversion gathered data container
        self._convert_ = {'setup' : [], 'solvation': []}
        self._convertdata_= None
        # Script part
        self._script_ = {'mainpage' : [], 'advanced': [], 'restrain': []}

        self.createmainPennon()

    def createmainPennon(self):
        '''Self explanatory neated with subroutines to make it more readable'''

        row = Frame(self,bg = "white")
        Label( row, bg = "white",
              image = self.img['logo']).pack( side= 'left', padx=25)
        row.pack(side="top", fill='x', padx=1)

        self.swapbody(1)

    def swapbody(self, _pbody_):# checked ok 16/09 -----------WF
        ''' Deletes and clean the last generated body
            maybe lacks a real body destroyer?? but works fine with
            this, because it is just a "small" overlapping I gess
        '''

        if self.prevailing_body <> _pbody_:
            if self.body == None:
                self.body = self.create_conversion_gui()

            else:
                self.body.destroy()

                if _pbody_==1:
                    print 'Swapping to gro2lam converter GUI'
                    self.body = self.create_conversion_gui()

                elif _pbody_==2:
                    print 'Swapping to input script generator GUI'
                    self.body = self.create_script_gui()

                elif _pbody_==3:
                    print 'Swapping to run script GUI'
                    self.body = self.create_run_gui()

                else:
                    exit('Wuut...')

            self.prevailing_body = _pbody_

            self.body.createWidgets()
            self.body.b1.focus()
            self.master.bind('<Return>', self.b1_hook )
            self.master.bind('<Escape>', self.quit_hook )
            self.body.pack(side='top', fill='x')

    def b1_hook(self, event=None):
        self.body.b1.invoke()

    def quit_hook(self, event=None):
        self.body.quit()

    def swap_hook(self):
        _l_ = [1,2,3]
        b = _l_[_l_.index(self.prevailing_body)-2]
        self.swapbody(b)

    def create_conversion_gui(self):
        'Hook to create conversion gui'
        return Conversion(self)# Hook

    def create_script_gui(self):
        'Hook to create script gui'
        return Script_GUI(self)# Hook

    def create_run_gui(self):
        'Hook to create run gui'
        return Run_GUI(self)# Hook

#------------------------------------------------------
'''///////////////     Sub routines     /////////////'''
#------------------------------------------------------

def launch_gui( started = False):
    ''' launcher
        Main GUI constructor
    '''

    print wrg_3('Before you start, make sure there are no comments',
                '(;) in the middle of a line of the input GROMACS files.',
                'Data after this symbol are not taken into account.')

    MasterWin = Tk()
    prompt = Gro2Lam_GUI( master= MasterWin, test = started)# xl_App

    # Top main pennon menu bar definition

    entry_list_of_dicts = [{ 'title' : 'File',
                            'cascade' : (('Quit' ,MasterWin.quit), ) },
                           { 'title' : 'Data File Creation',
                            'title_com' : (prompt.swapbody , 1)},
                           { 'title' : 'Input File Creation',
                            'title_com' : (prompt.swapbody , 2)},
                           { 'title' : 'Run',
                            'title_com' : (prompt.swapbody , 3)},
                           { 'titlei' : prompt.img['help'],
                            'cascade' : (('User manual' , showuserman),
                                         ('About' , launch_about, prompt),)}
                          ]
    createmenubar(MasterWin, entry_list_of_dicts)

    w = 460
    h = 570
    # get screen width and height
    ws = MasterWin.winfo_screenwidth() # width of the screen
    hs = MasterWin.winfo_screenheight() # height of the screen
    # calculate x and y coordinates for the Tk root window
    x = (ws/6) - (w/2)
    if x <100:
        x = 100
    y = (hs/3) - (h/2)
    if y< 40:
        y = 40

    prompt.MAINVERTEX = [ws, hs, w, h, x, y]
    #print MAINVERTEX
    # set the dimensions of the screen
    # and where it is placed
    MasterWin.geometry('{:d}x{:d}+{:d}+{:d}'.format( *prompt.MAINVERTEX[2:]))

    prompt.mainloop()

    try:
        MasterWin.destroy()
    except TclError:
        pass

def showlicence():

    print 'Opening licence file'
    command = 'gedit ./lib/docs/COPYING'#
    run_command(command)

def launch_about( _master_window_):

    print 'Launching about'

    title_txt = ' '*17+'ABOUT GROTOLAM'

    pop = AboutPopUp(master = _master_window_,
                     title = title_txt,
                     licence = showlicence
                    )

def showuserman():

    print 'Opening readme file'
    command = 'gedit ./lib/docs/README.md'#
    run_command(command)

# vim:tw=80


"""
Phong Material

For phong shading
"""
from .material import Material
from ..math import Vec3, Ray, HitRecord, dot3, reflect3, normalize3, clamp3
from ..camera import Camera

class PhongMaterial(Material):

    """Base Material Class"""

    def __init__(self, color: Vec3 = Vec3(1.,1.,1.), shininess: float = 10.0, reflectivity: float = 0.0, refraction: float = 1.0):
        Material.__init__(self, color, shininess, reflectivity, refraction)

    def shade(self, camera: Camera, ray: Ray, hitrecord: HitRecord,  lights: list) -> Vec3:
        """
        Shade method: Phong

        phong shader
        """
        colorsum = Vec3(0.,0.,0.)

        if len(lights)>0:
            for light in lights:
                N = hitrecord.normal_g
                L = normalize3(hitrecord.point - light.position)
                E = normalize3(camera.position - hitrecord.point)
                R = normalize3(-reflect3(L, N))
                diffuse = max(1. - dot3(N, L), 0.0)
                specular = pow(max(dot3(R, E), 0.0), 0.3 * self.shininess)
                color = self.color * 0.5 * (diffuse + specular) * hitrecord.color
                colorsum += color
            colorsum /= len(lights)
            colorsum = clamp3(colorsum, Vec3(0.,0.,0.), Vec3(1.,1.,1.))
        else:
            # no light in scene, use material color
            colorsum = self.color * hitrecord.color

        return colorsum


import numpy as np
import os
import pickle
import tensorflow as tf
import matplotlib.pyplot as plt
from skimage.transform import rotate, resize
from skimage import exposure
import skimage.io as io

from config import FLAGS

def load_facegreyreduxshuffled_set(batch_size, is_training=True):
    path = os.path.join('data', 'facegreyredux')
    if is_training:
        fd = open(os.path.join(path, 'facegreyredux'), 'rb')
        # loaded = np.fromfile(file=fd, dtype=np.uint8)
        loaded = np.asarray(pickle.load(fd))
        trainX = loaded.reshape((50000, 28, 28, 1)).astype(np.float32)

        fd = open(os.path.join(path, 'facegreyreduxcat'), 'rb')
        # loaded = np.fromfile(file=fd, dtype=np.uint8)
        loaded = np.asarray(pickle.load(fd))
        trainY = loaded.reshape((50000)).astype(np.int32)

        data_set = list(zip(trainX,trainY))
        np.random.shuffle(data_set)
        trainX, trainY = list(zip(*data_set))
        trainX = np.asarray(trainX).reshape((50000, 28, 28, 1)).astype(np.float32)
        trainY = np.asarray(trainY).reshape((50000)).astype(np.int32)
        trX = trainX[:40000] / 255.
        trY = trainY[:40000]

        valX = trainX[40000:, ] / 255.
        valY = trainY[40000:]

        num_tr_batch = 40000 // batch_size
        num_val_batch = 10000 // batch_size

        return trX, trY, num_tr_batch, valX, valY, num_val_batch
    else:
        if (FLAGS.flickr):
            fd = open(os.path.join(path, 'flickrsetgreyredux'), 'rb')
            loaded = np.asarray(pickle.load(fd))
            trainX = loaded.reshape((10000, 28, 28)).astype(np.float32) / 255.
        else:
            fd = open(os.path.join(path, 'facegreyreduxeval'), 'rb')
            loaded = np.asarray(pickle.load(fd))
            trainX = loaded.reshape((10000, 28, 28)).astype(np.float32) / 255.

        fd = open(os.path.join(path, 'facegreyreduxevalcat'), 'rb')
        loaded = np.asarray(pickle.load(fd))
        trainY = loaded.reshape((10000)).astype(np.int32)

        rotatedlist = []
        for image in trainX:
            image = rotate(image, FLAGS.rotation, preserve_range=True)
            if(FLAGS.mooney):
                v_min, v_max = np.percentile(image, (49.99999999, 51))
                image = exposure.rescale_intensity(image, in_range=(v_min, v_max))
            rotatedlist.append(image)
            if(len(rotatedlist)==1000):
                I = resize(image.reshape(28, 28), (128, 128))
                io.imsave("rotate" + str(FLAGS.rotation) +  "example.jpg", I, cmap='gray')
        rotatedlist = np.asarray(rotatedlist)
        plt.imshow(rotatedlist[33], cmap='gray')
        plt.show()
        trainX = rotatedlist.reshape((10000, 28, 28, 1)).astype(np.float32)

        return trainX, trainY

def create_inputs_norb(path, is_train: bool):
    """Get a batch from the input pipeline.

    Author:
      Ashley Gritzman 15/11/2018
    Args:
      is_train:
    Returns:
      img, lab:
    """
    if is_train:
        trX, trY, num_tr_batch, valX, valY, num_val_batch = load_facegreyreduxshuffled_set(FLAGS.batch_size, is_train)
    else:
        trX, trY = load_facegreyreduxshuffled_set(FLAGS.batch_size, is_train)

    def generator():
        for e1, e2 in zip(trX, trY):
            yield e1, e2

    capacity = 2000 + 3 * FLAGS.batch_size
    # Create batched dataset
    tf_dataset = tf.data.Dataset.from_generator(generator, output_types=(tf.float32, tf.int32), output_shapes=(tf.TensorShape(list(trX[0].shape)), ())).repeat().shuffle(capacity).batch(batch_size=FLAGS.batch_size, drop_remainder=True).prefetch(1)

    # dataset = input_fn(path, is_train)

    # Create one-shot iterator
    iterator = tf.compat.v1.data.make_one_shot_iterator(tf_dataset)

    img, lab = iterator.get_next()

    output_dict = {'image': img,
                   'label': lab}

    return output_dict


"""The tests for the  MQTT binary sensor platform."""
import copy
from datetime import datetime, timedelta
import json
from unittest.mock import patch

import pytest

from homeassistant.components import binary_sensor
from homeassistant.const import (
    EVENT_STATE_CHANGED,
    STATE_OFF,
    STATE_ON,
    STATE_UNAVAILABLE,
    STATE_UNKNOWN,
)
import homeassistant.core as ha
from homeassistant.setup import async_setup_component
import homeassistant.util.dt as dt_util

from .test_common import (
    help_test_availability_when_connection_lost,
    help_test_availability_without_topic,
    help_test_custom_availability_payload,
    help_test_default_availability_payload,
    help_test_discovery_broken,
    help_test_discovery_removal,
    help_test_discovery_update,
    help_test_discovery_update_attr,
    help_test_discovery_update_unchanged,
    help_test_entity_debug_info_message,
    help_test_entity_device_info_remove,
    help_test_entity_device_info_update,
    help_test_entity_device_info_with_connection,
    help_test_entity_device_info_with_identifier,
    help_test_entity_id_update_discovery_update,
    help_test_entity_id_update_subscriptions,
    help_test_setting_attribute_via_mqtt_json_message,
    help_test_setting_attribute_with_template,
    help_test_unique_id,
    help_test_update_with_json_attrs_bad_JSON,
    help_test_update_with_json_attrs_not_dict,
)

from tests.common import async_fire_mqtt_message, async_fire_time_changed

DEFAULT_CONFIG = {
    binary_sensor.DOMAIN: {
        "platform": "mqtt",
        "name": "test",
        "state_topic": "test-topic",
    }
}

async def test_setting_sensor_value_expires_availability_topic(hass, mqtt_mock, caplog):
    """Test the expiration of the value."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "expire_after": 4,
                "force_update": True,
                "availability_topic": "availability-topic",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

    async_fire_mqtt_message(hass, "availability-topic", "online")

    # State should be unavailable since expire_after is defined and > 0
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

    await expires_helper(hass, mqtt_mock, caplog)

async def test_setting_sensor_value_expires(hass, mqtt_mock, caplog):
    """Test the expiration of the value."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "expire_after": 4,
                "force_update": True,
            }
        },
    )
    await hass.async_block_till_done()

    # State should be unavailable since expire_after is defined and > 0
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

    await expires_helper(hass, mqtt_mock, caplog)

async def expires_helper(hass, mqtt_mock, caplog):
    """Run the basic expiry code."""
    realnow = dt_util.utcnow()
    now = datetime(realnow.year + 1, 1, 1, 1, tzinfo=dt_util.UTC)
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_time_changed(hass, now)
        async_fire_mqtt_message(hass, "test-topic", "ON")
        await hass.async_block_till_done()

    # Value was set correctly.
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    # Time jump +3s
    now = now + timedelta(seconds=3)
    async_fire_time_changed(hass, now)
    await hass.async_block_till_done()

    # Value is not yet expired
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    # Next message resets timer
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_time_changed(hass, now)
        async_fire_mqtt_message(hass, "test-topic", "OFF")
        await hass.async_block_till_done()

    # Value was updated correctly.
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF

    # Time jump +3s
    now = now + timedelta(seconds=3)
    async_fire_time_changed(hass, now)
    await hass.async_block_till_done()

    # Value is not yet expired
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF

    # Time jump +2s
    now = now + timedelta(seconds=2)
    async_fire_time_changed(hass, now)
    await hass.async_block_till_done()

    # Value is expired now
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

async def test_expiration_on_discovery_and_discovery_update_of_binary_sensor(
    hass, mqtt_mock, caplog
):
    """Test that binary_sensor with expire_after set behaves correctly on discovery and discovery update."""
    config = {
        "name": "Test",
        "state_topic": "test-topic",
        "expire_after": 4,
        "force_update": True,
    }

    config_msg = json.dumps(config)

    # Set time and publish config message to create binary_sensor via discovery with 4 s expiry
    realnow = dt_util.utcnow()
    now = datetime(realnow.year + 1, 1, 1, 1, tzinfo=dt_util.UTC)
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_time_changed(hass, now)
        async_fire_mqtt_message(
            hass, "homeassistant/binary_sensor/bla/config", config_msg
        )
        await hass.async_block_till_done()

    # Test that binary_sensor is not available
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

    # Publish state message
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_mqtt_message(hass, "test-topic", "ON")
        await hass.async_block_till_done()

    # Test that binary_sensor has correct state
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    # Advance +3 seconds
    now = now + timedelta(seconds=3)
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_time_changed(hass, now)
        await hass.async_block_till_done()

    # binary_sensor is not yet expired
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    # Resend config message to update discovery
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_time_changed(hass, now)
        async_fire_mqtt_message(
            hass, "homeassistant/binary_sensor/bla/config", config_msg
        )
        await hass.async_block_till_done()

    # Test that binary_sensor has not expired
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    # Add +2 seconds
    now = now + timedelta(seconds=2)
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_time_changed(hass, now)
        await hass.async_block_till_done()

    # Test that binary_sensor has expired
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

    # Resend config message to update discovery
    with patch(("homeassistant.helpers.event.dt_util.utcnow"), return_value=now):
        async_fire_mqtt_message(
            hass, "homeassistant/binary_sensor/bla/config", config_msg
        )
        await hass.async_block_till_done()

    # Test that binary_sensor is still expired
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNAVAILABLE

async def test_setting_sensor_value_via_mqtt_message(hass, mqtt_mock):
    """Test the setting of the value via MQTT."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")

    assert state.state == STATE_UNKNOWN

    async_fire_mqtt_message(hass, "test-topic", "ON")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    async_fire_mqtt_message(hass, "test-topic", "OFF")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF

async def test_invalid_sensor_value_via_mqtt_message(hass, mqtt_mock, caplog):
    """Test the setting of the value via MQTT."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")

    assert state.state == STATE_UNKNOWN

    async_fire_mqtt_message(hass, "test-topic", "0N")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNKNOWN
    assert "No matching payload found for entity" in caplog.text
    caplog.clear()
    assert "No matching payload found for entity" not in caplog.text

    async_fire_mqtt_message(hass, "test-topic", "ON")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    async_fire_mqtt_message(hass, "test-topic", "0FF")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON
    assert "No matching payload found for entity" in caplog.text

async def test_setting_sensor_value_via_mqtt_message_and_template(hass, mqtt_mock):
    """Test the setting of the value via MQTT."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
                "value_template": '{%if is_state(entity_id,"on")-%}OFF'
                "{%-else-%}ON{%-endif%}",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNKNOWN

    async_fire_mqtt_message(hass, "test-topic", "")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    async_fire_mqtt_message(hass, "test-topic", "")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF

async def test_setting_sensor_value_via_mqtt_message_and_template2(
    hass, mqtt_mock, caplog
):
    """Test the setting of the value via MQTT."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
                "value_template": "{{value | upper}}",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNKNOWN

    async_fire_mqtt_message(hass, "test-topic", "on")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    async_fire_mqtt_message(hass, "test-topic", "off")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF

    async_fire_mqtt_message(hass, "test-topic", "illegal")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF
    assert "template output: 'ILLEGAL'" in caplog.text

async def test_setting_sensor_value_via_mqtt_message_and_template_and_raw_state_encoding(
    hass, mqtt_mock, caplog
):
    """Test processing a raw value via MQTT."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "encoding": "",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
                "value_template": "{%if value|unpack('b')-%}ON{%else%}OFF{%-endif-%}",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNKNOWN

    async_fire_mqtt_message(hass, "test-topic", b"\x01")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

    async_fire_mqtt_message(hass, "test-topic", b"\x00")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF

async def test_setting_sensor_value_via_mqtt_message_empty_template(
    hass, mqtt_mock, caplog
):
    """Test the setting of the value via MQTT."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
                "value_template": '{%if value == "ABC"%}ON{%endif%}',
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNKNOWN

    async_fire_mqtt_message(hass, "test-topic", "DEF")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_UNKNOWN
    assert "Empty template output" in caplog.text

    async_fire_mqtt_message(hass, "test-topic", "ABC")
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON

async def test_valid_device_class(hass, mqtt_mock):
    """Test the setting of a valid sensor class."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "device_class": "motion",
                "state_topic": "test-topic",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state.attributes.get("device_class") == "motion"

async def test_invalid_device_class(hass, mqtt_mock):
    """Test the setting of an invalid sensor class."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "device_class": "abc123",
                "state_topic": "test-topic",
            }
        },
    )
    await hass.async_block_till_done()

    state = hass.states.get("binary_sensor.test")
    assert state is None

async def test_availability_when_connection_lost(hass, mqtt_mock):
    """Test availability after MQTT disconnection."""
    await help_test_availability_when_connection_lost(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_availability_without_topic(hass, mqtt_mock):
    """Test availability without defined availability topic."""
    await help_test_availability_without_topic(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_default_availability_payload(hass, mqtt_mock):
    """Test availability by default payload with defined topic."""
    await help_test_default_availability_payload(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_custom_availability_payload(hass, mqtt_mock):
    """Test availability by custom payload with defined topic."""
    await help_test_custom_availability_payload(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_force_update_disabled(hass, mqtt_mock):
    """Test force update option."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
            }
        },
    )
    await hass.async_block_till_done()

    events = []

    @ha.callback
    def callback(event):
        """Verify event got called."""
        events.append(event)

    hass.bus.async_listen(EVENT_STATE_CHANGED, callback)

    async_fire_mqtt_message(hass, "test-topic", "ON")
    await hass.async_block_till_done()
    assert len(events) == 1

    async_fire_mqtt_message(hass, "test-topic", "ON")
    await hass.async_block_till_done()
    assert len(events) == 1

async def test_force_update_enabled(hass, mqtt_mock):
    """Test force update option."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
                "force_update": True,
            }
        },
    )
    await hass.async_block_till_done()

    events = []

    @ha.callback
    def callback(event):
        """Verify event got called."""
        events.append(event)

    hass.bus.async_listen(EVENT_STATE_CHANGED, callback)

    async_fire_mqtt_message(hass, "test-topic", "ON")
    await hass.async_block_till_done()
    assert len(events) == 1

    async_fire_mqtt_message(hass, "test-topic", "ON")
    await hass.async_block_till_done()
    assert len(events) == 2

async def test_off_delay(hass, mqtt_mock):
    """Test off_delay option."""
    assert await async_setup_component(
        hass,
        binary_sensor.DOMAIN,
        {
            binary_sensor.DOMAIN: {
                "platform": "mqtt",
                "name": "test",
                "state_topic": "test-topic",
                "payload_on": "ON",
                "payload_off": "OFF",
                "off_delay": 30,
                "force_update": True,
            }
        },
    )
    await hass.async_block_till_done()

    events = []

    @ha.callback
    def callback(event):
        """Verify event got called."""
        events.append(event)

    hass.bus.async_listen(EVENT_STATE_CHANGED, callback)

    async_fire_mqtt_message(hass, "test-topic", "ON")
    await hass.async_block_till_done()
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON
    assert len(events) == 1

    async_fire_mqtt_message(hass, "test-topic", "ON")
    await hass.async_block_till_done()
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_ON
    assert len(events) == 2

    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=30))
    await hass.async_block_till_done()
    state = hass.states.get("binary_sensor.test")
    assert state.state == STATE_OFF
    assert len(events) == 3

async def test_setting_attribute_via_mqtt_json_message(hass, mqtt_mock):
    """Test the setting of attribute via MQTT with JSON payload."""
    await help_test_setting_attribute_via_mqtt_json_message(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_setting_attribute_with_template(hass, mqtt_mock):
    """Test the setting of attribute via MQTT with JSON payload."""
    await help_test_setting_attribute_with_template(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_update_with_json_attrs_not_dict(hass, mqtt_mock, caplog):
    """Test attributes get extracted from a JSON result."""
    await help_test_update_with_json_attrs_not_dict(
        hass, mqtt_mock, caplog, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_update_with_json_attrs_bad_JSON(hass, mqtt_mock, caplog):
    """Test attributes get extracted from a JSON result."""
    await help_test_update_with_json_attrs_bad_JSON(
        hass, mqtt_mock, caplog, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_discovery_update_attr(hass, mqtt_mock, caplog):
    """Test update of discovered MQTTAttributes."""
    await help_test_discovery_update_attr(
        hass, mqtt_mock, caplog, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_unique_id(hass, mqtt_mock):
    """Test unique id option only creates one sensor per unique_id."""
    config = {
        binary_sensor.DOMAIN: [
            {
                "platform": "mqtt",
                "name": "Test 1",
                "state_topic": "test-topic",
                "unique_id": "TOTALLY_UNIQUE",
            },
            {
                "platform": "mqtt",
                "name": "Test 2",
                "state_topic": "test-topic",
                "unique_id": "TOTALLY_UNIQUE",
            },
        ]
    }
    await help_test_unique_id(hass, mqtt_mock, binary_sensor.DOMAIN, config)

async def test_discovery_removal_binary_sensor(hass, mqtt_mock, caplog):
    """Test removal of discovered binary_sensor."""
    data = json.dumps(DEFAULT_CONFIG[binary_sensor.DOMAIN])
    await help_test_discovery_removal(
        hass, mqtt_mock, caplog, binary_sensor.DOMAIN, data
    )

async def test_discovery_update_binary_sensor_topic_template(hass, mqtt_mock, caplog):
    """Test update of discovered binary_sensor."""
    config1 = copy.deepcopy(DEFAULT_CONFIG[binary_sensor.DOMAIN])
    config2 = copy.deepcopy(DEFAULT_CONFIG[binary_sensor.DOMAIN])
    config1["name"] = "Beer"
    config2["name"] = "Milk"
    config1["state_topic"] = "sensor/state1"
    config2["state_topic"] = "sensor/state2"
    config1["value_template"] = "{{ value_json.state1.state }}"
    config2["value_template"] = "{{ value_json.state2.state }}"

    state_data1 = [
        ([("sensor/state1", '{"state1":{"state":"ON"}}')], "on", None),
    ]
    state_data2 = [
        ([("sensor/state2", '{"state2":{"state":"OFF"}}')], "off", None),
        ([("sensor/state2", '{"state2":{"state":"ON"}}')], "on", None),
        ([("sensor/state1", '{"state1":{"state":"OFF"}}')], "on", None),
        ([("sensor/state1", '{"state2":{"state":"OFF"}}')], "on", None),
        ([("sensor/state2", '{"state1":{"state":"OFF"}}')], "on", None),
        ([("sensor/state2", '{"state2":{"state":"OFF"}}')], "off", None),
    ]

    await help_test_discovery_update(
        hass,
        mqtt_mock,
        caplog,
        binary_sensor.DOMAIN,
        config1,
        config2,
        state_data1=state_data1,
        state_data2=state_data2,
    )

async def test_discovery_update_binary_sensor_template(hass, mqtt_mock, caplog):
    """Test update of discovered binary_sensor."""
    config1 = copy.deepcopy(DEFAULT_CONFIG[binary_sensor.DOMAIN])
    config2 = copy.deepcopy(DEFAULT_CONFIG[binary_sensor.DOMAIN])
    config1["name"] = "Beer"
    config2["name"] = "Milk"
    config1["state_topic"] = "sensor/state1"
    config2["state_topic"] = "sensor/state1"
    config1["value_template"] = "{{ value_json.state1.state }}"
    config2["value_template"] = "{{ value_json.state2.state }}"

    state_data1 = [
        ([("sensor/state1", '{"state1":{"state":"ON"}}')], "on", None),
    ]
    state_data2 = [
        ([("sensor/state1", '{"state2":{"state":"OFF"}}')], "off", None),
        ([("sensor/state1", '{"state2":{"state":"ON"}}')], "on", None),
        ([("sensor/state1", '{"state1":{"state":"OFF"}}')], "on", None),
        ([("sensor/state1", '{"state2":{"state":"OFF"}}')], "off", None),
    ]

    await help_test_discovery_update(
        hass,
        mqtt_mock,
        caplog,
        binary_sensor.DOMAIN,
        config1,
        config2,
        state_data1=state_data1,
        state_data2=state_data2,
    )

async def test_discovery_update_unchanged_binary_sensor(hass, mqtt_mock, caplog):
    """Test update of discovered binary_sensor."""
    config1 = copy.deepcopy(DEFAULT_CONFIG[binary_sensor.DOMAIN])
    config1["name"] = "Beer"

    data1 = json.dumps(config1)
    with patch(
        "homeassistant.components.mqtt.binary_sensor.MqttBinarySensor.discovery_update"
    ) as discovery_update:
        await help_test_discovery_update_unchanged(
            hass, mqtt_mock, caplog, binary_sensor.DOMAIN, data1, discovery_update
        )

@pytest.mark.no_fail_on_log_exception
async def test_discovery_broken(hass, mqtt_mock, caplog):
    """Test handling of bad discovery message."""
    data1 = '{ "name": "Beer",' '  "off_delay": -1 }'
    data2 = '{ "name": "Milk",' '  "state_topic": "test_topic" }'
    await help_test_discovery_broken(
        hass, mqtt_mock, caplog, binary_sensor.DOMAIN, data1, data2
    )

async def test_entity_device_info_with_connection(hass, mqtt_mock):
    """Test MQTT binary sensor device registry integration."""
    await help_test_entity_device_info_with_connection(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_entity_device_info_with_identifier(hass, mqtt_mock):
    """Test MQTT binary sensor device registry integration."""
    await help_test_entity_device_info_with_identifier(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_entity_device_info_update(hass, mqtt_mock):
    """Test device registry update."""
    await help_test_entity_device_info_update(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_entity_device_info_remove(hass, mqtt_mock):
    """Test device registry remove."""
    await help_test_entity_device_info_remove(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_entity_id_update_subscriptions(hass, mqtt_mock):
    """Test MQTT subscriptions are managed when entity_id is updated."""
    await help_test_entity_id_update_subscriptions(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_entity_id_update_discovery_update(hass, mqtt_mock):
    """Test MQTT discovery update when entity_id is updated."""
    await help_test_entity_id_update_discovery_update(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )

async def test_entity_debug_info_message(hass, mqtt_mock):
    """Test MQTT debug info."""
    await help_test_entity_debug_info_message(
        hass, mqtt_mock, binary_sensor.DOMAIN, DEFAULT_CONFIG
    )


import csv
from collections import OrderedDict
from datetime import datetime
from pathlib import Path
from typing import Any, List, Mapping

from dmutils.formats import DATE_FORMAT, DATETIME_FORMAT
from dmutils.s3 import S3

from dmscripts.helpers.s3_helpers import get_bucket_name

# This URL is framework agnostic
PUBLIC_BRIEF_URL = "https://www.digitalmarketplace.service.gov.uk/digital-outcomes-and-specialists/opportunities/{}"

DOS_OPPORTUNITY_HEADERS = [
    "ID", "Opportunity", "Link", "Framework", "Category", "Specialist",
    "Organisation Name", "Buyer Domain", "Location Of The Work",
    "Published At", "Open For", "Expected Contract Length", "Applications from SMEs",
    "Applications from Large Organisations", "Total Organisations", "Status", "Winning supplier",
    "Size of supplier", "Contract amount", "Contract start date", "Clarification questions", "Employment status"
]

DOWNLOAD_FILE_NAME = "opportunity-data.csv"

def format_datetime_string_as_date(dt):
    return datetime.strptime(dt, DATETIME_FORMAT).strftime(DATE_FORMAT) if dt else None

def remove_username_from_email_address(ea):
    return '{}'.format(ea.split('@').pop()) if ea else None

def _build_row(
    brief: dict, brief_responses: List[dict], include_buyer_user_details: bool = False
) -> OrderedDict:
    winner = None
    applications_from_sme_suppliers = 0
    applications_from_large_suppliers = 0

    for brief_response in brief_responses:
        if brief_response['supplierOrganisationSize'] == 'large':
            applications_from_large_suppliers += 1
        else:
            applications_from_sme_suppliers += 1

        if brief_response['status'] == 'awarded':
            winner = brief_response

    row = OrderedDict(zip(DOS_OPPORTUNITY_HEADERS, [
        brief['id'],
        brief['title'],
        PUBLIC_BRIEF_URL.format(brief['id']),
        brief['frameworkSlug'],
        brief['lotSlug'],
        brief.get('specialistRole', ""),
        brief['organisation'],
        remove_username_from_email_address(brief['users'][0]['emailAddress']),
        brief['location'],
        format_datetime_string_as_date(brief['publishedAt']),
        brief.get('requirementsLength', '2 weeks'),  # only briefs on the specialists lot include 'requirementsLength'
        brief.get('contractLength', ''),
        applications_from_sme_suppliers,
        applications_from_large_suppliers,
        applications_from_sme_suppliers + applications_from_large_suppliers,
        brief['status'],
        winner['supplierName'] if winner else '',
        winner['supplierOrganisationSize'] if winner else '',
        winner['awardDetails']['awardedContractValue'] if winner else '',
        winner['awardDetails']['awardedContractStartDate'] if winner else '',
        len(brief['clarificationQuestions']),
        brief.get('employmentStatus', ''),
    ]))

    if include_buyer_user_details:
        buyer_user = brief["users"][0]
        row.update([
            ("Buyer user name", buyer_user["name"]),
            ("Buyer email address", buyer_user["emailAddress"]),
            ("Buyer phone number", buyer_user.get("phoneNumber", "")),
        ])

    return row

def get_latest_dos_framework(client) -> str:
    frameworks = client.find_frameworks()['frameworks']
    for framework in frameworks:
        # Should be maximum of 1 live DOS framework
        if framework['family'] == 'digital-outcomes-and-specialists' and framework['status'] == 'live':
            return framework['slug']
    return 'digital-outcomes-and-specialists'

def get_brief_data(client, logger, include_buyer_user_details: bool = False) -> list:
    logger.info("Fetching closed briefs from API")
    briefs = client.find_briefs_iter(status="closed,awarded,unsuccessful,cancelled", with_users=True,
                                     with_clarification_questions=True)
    rows = []
    for brief in briefs:
        logger.info(f"Fetching brief responses for Brief ID {brief['id']}")
        brief_responses = client.find_brief_responses_iter(brief_id=brief['id'])
        rows.append(_build_row(brief, brief_responses, include_buyer_user_details))
    return rows

def write_rows_to_csv(rows: List[Mapping[str, Any]], file_path: Path, logger) -> None:
    logger.info(f"Writing rows to {file_path}")

    # assumes all rows have the same keys
    fieldnames = list(rows[0].keys())

    with open(file_path, 'w') as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames, delimiter=',', quotechar='"')
        writer.writeheader()
        for row in rows:
            writer.writerow(row)

def upload_file_to_s3(
    file_path,
    bucket,
    remote_key_name: str,
    download_name: str,
    *,
    public: bool = True,
    dry_run: bool = False,
    logger,
):
    with open(file_path, 'br') as source_file:
        acl = "public-read" if public else "bucket-owner-full-control"

        logger.info("{}UPLOAD: {} to s3://{}/{} with acl {}".format(
            '[Dry-run]' if dry_run else '',
            file_path,
            bucket.bucket_name,
            remote_key_name,
            acl
        ))

        if not dry_run:
            # Save file
            bucket.save(
                remote_key_name,
                source_file,
                acl=acl,
                download_filename=download_name
            )

def export_dos_opportunities(
    client,
    logger,
    stage: str,
    output_dir,
    dry_run: bool = False
):
    output_dir = Path(output_dir)
    if not output_dir.exists():
        logger.info(f"Creating {output_dir} directory")
        output_dir.mkdir(parents=True)

    latest_framework_slug = get_latest_dos_framework(client)

    communications_bucket = S3(get_bucket_name(stage, "communications"))
    reports_bucket = S3(get_bucket_name(stage, "reports"))

    logger.info("Exporting DOS opportunity data to CSV")

    # Get the data
    rows = get_brief_data(client, logger, include_buyer_user_details=True)

    # Construct CSV for admins
    write_rows_to_csv(rows, output_dir / "opportunity-data-for-admins.csv", logger)
    # Construct public CSV (filter out buyer details)
    write_rows_to_csv(
        [
            OrderedDict((k, v) for k, v in row.items() if k in DOS_OPPORTUNITY_HEADERS)
            for row in rows
        ],
        output_dir / DOWNLOAD_FILE_NAME,
        logger
    )

    # Upload admin CSV to reports bucket
    upload_file_to_s3(
        output_dir / "opportunity-data-for-admins.csv",
        reports_bucket,
        f"{latest_framework_slug}/reports/{DOWNLOAD_FILE_NAME}",
        DOWNLOAD_FILE_NAME,
        public=False,
        dry_run=dry_run,
        logger=logger
    )

    # Upload public CSV to S3
    upload_file_to_s3(
        output_dir / DOWNLOAD_FILE_NAME,
        communications_bucket,
        f"{latest_framework_slug}/communications/data/{DOWNLOAD_FILE_NAME}",
        DOWNLOAD_FILE_NAME,
        public=True,
        dry_run=dry_run,
        logger=logger
    )


"""Utility functions."""
import logging

import numpy as np
from scipy.signal import periodogram

from tensorpac.methods.meth_pac import _kl_hr
from tensorpac.pac import _PacObj, _PacVisual
from tensorpac.io import set_log_level

from matplotlib.gridspec import GridSpec
import matplotlib.pyplot as plt

logger = logging.getLogger('tensorpac')

def pac_vec(f_pha='mres', f_amp='mres'):
    """Generate cross-frequency coupling vectors.

    Parameters
    ----------
    Frequency vector for the phase and amplitude. Here you can use
    several forms to define those vectors :

        * Basic list/tuple (ex: [2, 4] or [8, 12]...)
        * List of frequency bands (ex: [[2, 4], [5, 7]]...)
        * Dynamic definition : (start, stop, width, step)
        * Range definition (ex : np.arange(3) => [[0, 1], [1, 2]])
        * Using a string. `f_pha` and `f_amp` can be 'lres', 'mres', 'hres'
          respectively for low, middle and high resolution vectors. In that
          case, it uses the definition proposed by Bahramisharif et al. 2013
          :cite:`bahramisharif2013propagating` i.e
          f_pha = [f - f / 4, f + f / 4] and f_amp = [f - f / 8, f + f / 8]

    Returns
    -------
    f_pha, f_amp : array_like
        Arrays containing the pairs of phase and amplitude frequencies. Each
        vector have a shape of (N, 2).
    """
    nb_fcy = dict(lres=10, mres=30, hres=50, demon=70, hulk=100)
    if isinstance(f_pha, str):
        # get where phase frequencies start / finish / number
        f_pha_start, f_pha_end = 2, 20
        f_pha_nb = nb_fcy[f_pha]
        # f_pha = [f - f / 4, f + f / 4]
        f_pha_mid = np.linspace(f_pha_start, f_pha_end, f_pha_nb)
        f_pha = np.c_[f_pha_mid - f_pha_mid / 4., f_pha_mid + f_pha_mid / 4.]
    if isinstance(f_amp, str):
        # get where amplitude frequencies start / finish / number
        f_amp_start, f_amp_end = 60, 160
        f_amp_nb = nb_fcy[f_amp]
        # f_amp = [f - f / 8, f + f / 8]
        f_amp_mid = np.linspace(f_amp_start, f_amp_end, f_amp_nb)
        f_amp = np.c_[f_amp_mid - f_amp_mid / 8., f_amp_mid + f_amp_mid / 8.]

    return _check_freq(f_pha), _check_freq(f_amp)

def _check_freq(f):
    """Check the frequency definition."""
    f = np.atleast_2d(np.asarray(f))
    #
    if len(f.reshape(-1)) == 1:
        raise ValueError("The length of f should at least be 2.")
    elif 2 in f.shape:  # f of shape (N, 2) or (2, N)
        if f.shape[1] is not 2:
            f = f.T
    elif np.squeeze(f).shape == (4,):  # (f_start, f_end, f_width, f_step)
        f = _pair_vectors(*tuple(np.squeeze(f)))
    else:  # Sequential
        f = f.reshape(-1)
        f.sort()
        f = np.c_[f[0:-1], f[1::]]
    return f

def _pair_vectors(f_start, f_end, f_width, f_step):
    # Generate two array for phase and amplitude :
    fdown = np.arange(f_start, f_end - f_width, f_step)
    fup = np.arange(f_start + f_width, f_end, f_step)
    return np.c_[fdown, fup]

def pac_trivec(f_start=60., f_end=160., f_width=10.):
    """Generate triangular vector.

    By contrast with the pac_vec function, this function generate frequency
    vector with an increasing frequency bandwidth.

    Parameters
    ----------
    f_start : float | 60.
        Starting frequency.
    f_end : float | 160.
        Ending frequency.
    f_width : float | 10.
        Frequency bandwidth increase between each band.

    Returns
    -------
    f : array_like
        The triangular vector.
    tridx : array_like
        The triangular index for the reconstruction.
    """
    starting = np.arange(f_start, f_end + f_width, f_width)
    f, tridx = np.array([]), np.array([])
    for num, k in enumerate(starting[0:-1]):
        # Lentgh of the vector to build :
        le = len(starting) - (num + 1)
        # Create the frequency vector for this starting frequency :
        fst = np.c_[np.full(le, k), starting[num + 1::]]
        nfst = fst.shape[0]
        # Create the triangular index for this vector of frequencies :
        idx = np.c_[np.flipud(np.arange(nfst)), np.full(nfst, num)]
        tridx = np.concatenate((tridx, idx), axis=0) if tridx.size else idx
        f = np.concatenate((f, fst), axis=0) if f.size else fst
    return f, tridx

class PSD(object):
    """Power Spectrum Density for electrophysiological brain data.

    Parameters
    ----------
    x : array_like
        Array of data of shape (n_epochs, n_times)
    sf : float
        The sampling frequency.
    """

    def __init__(self, x, sf):
        """Init."""
        assert isinstance(x, np.ndarray) and (x.ndim == 2), (
            "x should be a 2d array of shape (n_epochs, n_times)")
        self._n_trials, self._n_times = x.shape
        logger.info(f"Compute PSD over {self._n_trials} trials and "
                    f"{self._n_times} time points")
        self._freqs, self._psd = periodogram(x, fs=sf, window=None,
                                             nfft=self._n_times,
                                             detrend='constant',
                                             return_onesided=True,
                                             scaling='density', axis=1)

    def plot(self, f_min=None, f_max=None, confidence=95, interp=None,
             log=False, grid=True, fz_title=18, fz_labels=15):
        """Plot the PSD.

        Parameters
        ----------
        f_min, f_max : (int, float) | None
            Frequency bounds to use for plotting
        confidence : (int, float) | None
            Light gray confidence interval. If None, no interval will be
            displayed
        interp : int | None
            Line interpolation integer. For example, if interp is 10 the number
            of points is going to be multiply by 10
        log : bool | False
            Use a log scale representation
        grid : bool | True
            Add a grid to the plot
        fz_title : int | 18
            Font size for the title
        fz_labels : int | 15
            Font size the x/y labels

        Returns
        -------
        ax : Matplotlib axis
            The matplotlib axis that contains the figure
        """
        import matplotlib.pyplot as plt
        f_types = (int, float)
        # interpolation
        xvec, yvec = self._freqs, self._psd
        if isinstance(interp, int) and (interp > 1):
            # from scipy.interpolate import make_interp_spline, BSpline
            from scipy.interpolate import interp1d
            xnew = np.linspace(xvec[0], xvec[-1], len(xvec) * interp)
            f = interp1d(xvec, yvec, kind='quadratic', axis=1)
            yvec = f(xnew)
            xvec = xnew
        # (f_min, f_max)
        f_min = xvec[0] if not isinstance(f_min, f_types) else f_min
        f_max = xvec[-1] if not isinstance(f_max, f_types) else f_max
        # plot main psd
        plt.plot(xvec, yvec.mean(0), color='black',
                 label='mean PSD over trials')
        # plot confidence interval
        if isinstance(confidence, (int, float)) and (0 < confidence < 100):
            logger.info(f"    Add {confidence}th confidence interval")
            interval = (100. - confidence) / 2
            kw = dict(axis=0, interpolation='nearest')
            psd_min = np.percentile(yvec, interval, **kw)
            psd_max = np.percentile(yvec, 100. - interval, **kw)
            plt.fill_between(xvec, psd_max, psd_min, color='lightgray',
                             alpha=0.5,
                             label=f"{confidence}th confidence interval")
            plt.legend(fontsize=fz_labels)
        plt.xlabel("Frequencies (Hz)", fontsize=fz_labels)
        plt.ylabel("Power (V**2/Hz)", fontsize=fz_labels)
        plt.title(f"PSD mean over {self._n_trials} trials", fontsize=fz_title)
        plt.xlim(f_min, f_max)
        if log:
            from matplotlib.ticker import ScalarFormatter
            plt.xscale('log', basex=10)
            plt.gca().xaxis.set_major_formatter(ScalarFormatter())
        if grid:
            plt.grid(color='grey', which='major', linestyle='-',
                     linewidth=1., alpha=0.5)
            plt.grid(color='lightgrey', which='minor', linestyle='--',
                     linewidth=0.5, alpha=0.5)

        return plt.gca()

    def plot_st_psd(self, f_min=None, f_max=None, log=False, grid=True,
                    fz_title=18, fz_labels=15, fz_cblabel=15, **kw):
        """Single-trial PSD plot.

        Parameters
        ----------
        f_min, f_max : (int, float) | None
            Frequency bounds to use for plotting
        log : bool | False
            Use a log scale representation
        grid : bool | True
            Add a grid to the plot
        fz_title : int | 18
            Font size for the title
        fz_labels : int | 15
            Font size the x/y labels
        fz_cblabel : int | 15
            Font size the colorbar label labels

        Returns
        -------
        ax : Matplotlib axis
            The matplotlib axis that contains the figure
        """
        # manage input variables
        kw['fz_labels'] = kw.get('fz_labels', fz_labels)
        kw['fz_title'] = kw.get('fz_title', fz_title)
        kw['fz_cblabel'] = kw.get('fz_cblabel', fz_title)
        kw['xlabel'] = kw.get('xlabel', "Frequencies (Hz)")
        kw['ylabel'] = kw.get('ylabel', "Trials")
        kw['title'] = kw.get('title', "Single-trial PSD")
        kw['cblabel'] = kw.get('cblabel', "Power (V**2/Hz)")
        # (f_min, f_max)
        xvec, psd = self._freqs, self._psd
        f_types = (int, float)
        f_min = xvec[0] if not isinstance(f_min, f_types) else f_min
        f_max = xvec[-1] if not isinstance(f_max, f_types) else f_max
        # locate (f_min, f_max) indices
        f_min_idx = np.abs(xvec - f_min).argmin()
        f_max_idx = np.abs(xvec - f_max).argmin()
        sl_freq = slice(f_min_idx, f_max_idx)
        xvec = xvec[sl_freq]
        psd = psd[:, sl_freq]
        # make the 2D plot
        _viz = _PacVisual()
        trials = np.arange(self._n_trials)
        _viz.pacplot(psd, xvec, trials, **kw)
        if log:
            from matplotlib.ticker import ScalarFormatter
            plt.xscale('log', basex=10)
            plt.gca().xaxis.set_major_formatter(ScalarFormatter())
        if grid:
            plt.grid(color='grey', which='major', linestyle='-',
                     linewidth=1., alpha=0.5)
            plt.grid(color='lightgrey', which='minor', linestyle='--',
                     linewidth=0.5, alpha=0.5)

        return plt.gca()

    def show(self):
        """Display the PSD figure."""
        import matplotlib.pyplot as plt
        plt.show()

    @property
    def freqs(self):
        """Get the frequency vector."""
        return self._freqs

    @property
    def psd(self):
        """Get the psd value."""
        return self._psd

class BinAmplitude(_PacObj):
    """Bin the amplitude according to the phase.

    Parameters
    ----------
    x : array_like
        Array of data of shape (n_epochs, n_times)
    sf : float
        The sampling frequency
    f_pha : tuple, list | [2, 4]
        List of two floats describing the frequency bounds for extracting the
        phase
    f_amp : tuple, list | [60, 80]
        List of two floats describing the frequency bounds for extracting the
        amplitude
    n_bins : int | 18
        Number of bins to use to binarize the phase and the amplitude
    dcomplex : {'wavelet', 'hilbert'}
        Method for the complex definition. Use either 'hilbert' or
        'wavelet'.
    cycle : tuple | (3, 6)
        Control the number of cycles for filtering (only if dcomplex is
        'hilbert'). Should be a tuple of integers where the first one
        refers to the number of cycles for the phase and the second for the
        amplitude :cite:`bahramisharif2013propagating`.
    width : int | 7
        Width of the Morlet's wavelet.
    edges : int | None
        Number of samples to discard to avoid edge effects due to filtering
    """

    def __init__(self, x, sf, f_pha=[2, 4], f_amp=[60, 80], n_bins=18,
                 dcomplex='hilbert', cycle=(3, 6), width=7, edges=None,
                 n_jobs=-1):
        """Init."""
        _PacObj.__init__(self, f_pha=f_pha, f_amp=f_amp, dcomplex=dcomplex,
                         cycle=cycle, width=width)
        # check
        x = np.atleast_2d(x)
        assert x.ndim <= 2, ("`x` input should be an array of shape "
                             "(n_epochs, n_times)")
        assert isinstance(sf, (int, float)), ("`sf` input should be a integer "
                                              "or a float")
        assert all([isinstance(k, (int, float)) for k in f_pha]), (
            "`f_pha` input should be a list of two integers / floats")
        assert all([isinstance(k, (int, float)) for k in f_amp]), (
            "`f_amp` input should be a list of two integers / floats")
        assert isinstance(n_bins, int), "`n_bins` should be an integer"
        logger.info(f"Binning {f_amp}Hz amplitude according to {f_pha}Hz "
                    "phase")
        # extract phase and amplitude
        kw = dict(keepfilt=False, edges=edges, n_jobs=n_jobs)
        pha = self.filter(sf, x, 'phase', **kw)
        amp = self.filter(sf, x, 'amplitude', **kw)
        # binarize amplitude according to phase
        self._amplitude = _kl_hr(pha, amp, n_bins, mean_bins=False).squeeze()
        self.n_bins = n_bins

    def plot(self, unit='rad', normalize=False, **kw):
        """Plot the amplitude.

        Parameters
        ----------
        unit : {'rad', 'deg'}
            The unit to use for the phase. Use either 'deg' for degree or 'rad'
            for radians
        normalize : bool | None
            Normalize the histogram by the maximum
        kw : dict | {}
            Additional inputs are passed to the matplotlib.pyplot.bar function

        Returns
        -------
        ax : Matplotlib axis
            The matplotlib axis that contains the figure
        """
        import matplotlib.pyplot as plt
        assert unit in ['rad', 'deg']
        if unit == 'rad':
            self._phase = np.linspace(-np.pi, np.pi, self.n_bins)
            width = 2 * np.pi / self.n_bins
        elif unit == 'deg':
            self._phase = np.linspace(-180, 180, self.n_bins)
            width = 360 / self.n_bins
        amp_mean = self._amplitude.mean(1)
        if normalize:
            amp_mean /= amp_mean.max()
        plt.bar(self._phase, amp_mean, width=width, **kw)
        plt.xlabel(f"Frequency phase ({self.n_bins} bins)", fontsize=18)
        plt.ylabel("Amplitude", fontsize=18)
        plt.title("Binned amplitude")
        plt.autoscale(enable=True, axis='x', tight=True)

    def show(self):
        """Show the figure."""
        import matplotlib.pyplot as plt
        plt.show()

    @property
    def amplitude(self):
        """Get the amplitude value."""
        return self._amplitude

    @property
    def phase(self):
        """Get the phase value."""
        return self._phase

class ITC(_PacObj, _PacVisual):
    """Compute the Inter-Trials Coherence (ITC).

    The Inter-Trials Coherence (ITC) is a measure of phase consistency over
    trials for a single recording site (electrode / sensor etc.).

    Parameters
    ----------
    x : array_like
        Array of data of shape (n_epochs, n_times)
    sf : float
        The sampling frequency
    f_pha : tuple, list | [2, 4]
        List of two floats describing the frequency bounds for extracting the
        phase
    dcomplex : {'wavelet', 'hilbert'}
        Method for the complex definition. Use either 'hilbert' or
        'wavelet'.
    cycle : tuple | 3
        Control the number of cycles for filtering the phase (only if dcomplex
        is 'hilbert').
    width : int | 7
        Width of the Morlet's wavelet.
    edges : int | None
        Number of samples to discard to avoid edge effects due to filtering
    """

    def __init__(self, x, sf, f_pha=[2, 4], dcomplex='hilbert', cycle=3,
                 width=7, edges=None, n_jobs=-1, verbose=None):
        """Init."""
        set_log_level(verbose)
        _PacObj.__init__(self, f_pha=f_pha, f_amp=[60, 80], dcomplex=dcomplex,
                         cycle=(cycle, 6), width=width)
        _PacVisual.__init__(self)
        # check
        x = np.atleast_2d(x)
        assert x.ndim <= 2, ("`x` input should be an array of shape "
                             "(n_epochs, n_times)")
        self._n_trials = x.shape[0]
        logger.info("Inter-Trials Coherence (ITC)")
        logger.info(f"    extracting {len(self.xvec)} phases")
        # extract phase and amplitude
        kw = dict(keepfilt=False, edges=edges, n_jobs=n_jobs)
        pha = self.filter(sf, x, 'phase', **kw)
        # compute itc
        self._itc = np.abs(np.exp(1j * pha).mean(1)).squeeze()
        self._sf = sf

    def plot(self, times=None, **kw):
        """Plot the Inter-Trials Coherence.

        Parameters
        ----------
        times : array_like | None
            Custom time vector to use
        kw : dict | {}
            Additional inputs are either pass to the matplotlib.pyplot.plot
            function if a single phase band is used, otherwise to the
            matplotlib.pyplot.pcolormesh function

        Returns
        -------
        ax : Matplotlib axis
            The matplotlib axis that contains the figure
        """
        import matplotlib.pyplot as plt
        n_pts = self._itc.shape[-1]
        if not isinstance(times, np.ndarray):
            times = np.arange(n_pts) / self._sf
        times = times[self._edges]
        assert len(times) == n_pts, ("The length of the time vector should be "
                                     "{n_pts}")
        xlab = 'Time'
        title = f"Inter-Trials Coherence ({self._n_trials} trials)"
        if self._itc.ndim == 1:
            plt.plot(times, self._itc, **kw)
        elif self._itc.ndim == 2:
            vmin = kw.get('vmin', np.percentile(self._itc, 1))
            vmax = kw.get('vmax', np.percentile(self._itc, 99))
            self.pacplot(self._itc, times, self.xvec, vmin=vmin, vmax=vmax,
                         ylabel="Frequency for phase (Hz)", xlabel=xlab,
                         title=title, **kw)
        return plt.gca()

    def show(self):
        """Show the figure."""
        import matplotlib.pyplot as plt
        plt.show()

    @property
    def itc(self):
        """Get the itc value."""
        return self._itc

class PeakLockedTF(_PacObj, _PacVisual):
    """Peak-Locked Time-frequency representation.

    This class can be used in order to re-align time-frequency representations
    around a time-point (cue) according to the closest phase peak. This type
    of visualization can bring out a cyclic behavior of the amplitude at a
    given phase, potentially indicating the presence of a phase-amplitude
    coupling. Here's the detailed pipeline :

        * Filter around a single phase frequency bands and across multiple
          amplitude frequencies
        * Use a `cue` which define the time-point to use for the realignment
        * Detect in the filtered phase the closest peak to the cue. This step
          is repeated to each trial in order to get a list of length (n_epochs)
          that contains the number of sample (shift) so that if the phase is
          moved, the peak fall onto the cue. A positive shift indicates that
          the phase is moved forward while a negative shift is for a backward
          move
        * Apply, to each trial, this shift to the amplitude
        * Plot the mean re-aligned amplitudes

    Parameters
    ----------
    x : array_like
        Array of data of shape (n_epochs, n_times)
    sf : float
        The sampling frequency
    cue : int, float
        Time-point to use in order to detect the closest phase peak. This
        parameter works in conjunction with the `times` input below. Use
        either :

            * An integer and `times` is None to indicate that you want to
              realign according to a time-point in sample
            * A integer or a float with `times` the time vector if you want
              that Tensorpac automatically infer the sample number around which
              to align
    times : array_like | None
        Time vector
    f_pha : tuple, list | [2, 4]
        List of two floats describing the frequency bounds for extracting the
        phase
    f_amp : tuple, list | [60, 80]
        Frequency vector for the amplitude. Here you can use several forms to
        define those vectors :

            * Dynamic definition : (start, stop, width, step)
            * Using a string : `f_amp` can be 'lres', 'mres', 'hres'
              respectively for low, middle and high resolution vectors
    cycle : tuple | (3, 6)
        Control the number of cycles for filtering. Should be a tuple of
        integers where the first one refers to the number of cycles for the
        phase and the second for the amplitude
        :cite:`bahramisharif2013propagating`.
    """

    def __init__(self, x, sf, cue, times=None, f_pha=[5, 7], f_amp='hres',
                 cycle=(3, 6), n_jobs=-1, verbose=None):
        """Init."""
        set_log_level(verbose)
        # initialize to retrieve filtering methods
        _PacObj.__init__(self, f_pha=f_pha, f_amp=f_amp, dcomplex='hilbert',
                         cycle=cycle)
        _PacVisual.__init__(self)
        logger.info("PeakLockedTF object defined")
        # inputs checking
        x = np.atleast_2d(x)
        assert isinstance(x, np.ndarray) and (x.ndim == 2)
        assert isinstance(sf, (int, float))
        assert isinstance(cue, (int, float))
        assert isinstance(f_pha, (list, tuple)) and (len(f_pha) == 2)
        n_epochs, n_times = x.shape

        # manage cur conversion
        if times is None:
            cue = int(cue)
            times = np.arange(n_times)
            logger.info(f"    align on sample cue={cue}")
        else:
            assert isinstance(times, np.ndarray) and (len(times) == n_times)
            cue_time = cue
            cue = np.abs(times - cue).argmin() - 1
            logger.info(f"    align on time-point={cue_time} (sample={cue})")
        self.cue, self._times = cue, times

        # extract phase and amplitudes
        logger.info(f"    extract phase and amplitudes "
                    f"(n_amps={len(self.yvec)})")
        kw = dict(keepfilt=False, n_jobs=n_jobs)
        pha = self.filter(sf, x, 'phase', n_jobs=n_jobs, keepfilt=True)
        amp = self.filter(sf, x, 'amplitude', n_jobs=n_jobs)
        self._pha, self._amp = pha, amp ** 2

        # peak detection
        logger.info(f"    running peak detection around sample={cue}")
        self.shifts = self._peak_detection(self._pha.squeeze(), cue)

        # realign phases and amplitudes
        logger.info(f"    realign the {n_epochs} phases and amplitudes")
        self.amp_a = self._shift_signals(self._amp, self.shifts, fill_with=0.)
        self.pha_a = self._shift_signals(self._pha, self.shifts, fill_with=0.)

    @staticmethod
    def _peak_detection(pha, cue):
        """Single trial closest to a cue peak detection.

        Parameters
        ----------
        pha : array_like
            Array of single trial phases of shape (n_trials, n_times)
        cue : int
            Cue to use as a reference (in sample unit)

        Returns
        -------
        peaks : array_like
            Array of length (n_trials,) describing each delay to apply
            to each trial in order to realign the phases. In detail :

                * Positive delays means that zeros should be prepend
                * Negative delays means that zeros should be append
        """
        n_trials, n_times = pha.shape
        peaks = []
        for tr in range(n_trials):
            # select the single trial phase
            st_pha = pha[tr, :]
            # detect all peaks across time points
            st_peaks = []
            for t in range(n_times - 1):
                if (st_pha[t - 1] < st_pha[t]) and (st_pha[t] > st_pha[t + 1]):
                    st_peaks += [t]
            # detect the minimum peak
            min_peak = st_peaks[np.abs(np.array(st_peaks) - cue).argmin()]
            peaks += [cue - min_peak]

        return np.array(peaks)

    @staticmethod
    def _shift_signals(sig, n_shifts, fill_with=0):
        """Shift an array of signals according to an array of delays.

        Parameters
        ----------
        sig : array_like
            Array of signals of shape (n_freq, n_trials, n_times)
        n_shifts : array_like
            Array of delays to apply to each trial of shape (n_trials,)
        fill_with : int
            Value to prepend / append to each shifted time-series

        Returns
        -------
        sig_shifted : array_like
            Array of shifted signals with the same shape as the input
        """
        # prepare the needed variables
        n_freqs, n_trials, n_pts = sig.shape
        sig_shifted = np.zeros_like(sig)
        # shift each trial
        for tr in range(n_trials):
            # select the data of a specific trial
            st_shift = n_shifts[tr]
            st_sig = sig[:, tr, :]
            fill = np.full((n_freqs, abs(st_shift)), fill_with,
                           dtype=st_sig.dtype)
            # shift this specific trial
            if st_shift > 0:   # move forward = prepend zeros
                sig_shifted[:, tr, :] = np.c_[fill, st_sig][:, 0:-st_shift]
            elif st_shift < 0:  # move backward = append zeros
                sig_shifted[:, tr, :] = np.c_[st_sig, fill][:, abs(st_shift):]

        return sig_shifted

    def plot(self, zscore=False, baseline=None, edges=0, **kwargs):
        """Integrated Peak-Locked TF plotting function.

        Parameters
        ----------
        zscore : bool | False
            Normalize the power by using a z-score normalization. This can be
            useful in order to compensate the 1 / f effect in the power
            spectrum. If True, the mean and deviation are computed at the
            single trial level and across all time points
        baseline : tuple | None
            Baseline period to use in order to apply the z-score correction.
            Should be in samples.
        edges : int | 0
            Number of pixels to discard to compensate filtering edge effect
            (`power[edges:-edges]`).
        kwargs : dict | {}
            Additional arguments are sent to the
            :class:`tensorpac.utils.PeakLockedTF.pacplot` method
        """
        # manage additional arguments
        kwargs['colorbar'] = False
        kwargs['ylabel'] = 'Frequency for amplitude (hz)'
        kwargs['xlabel'] = ''
        kwargs['fz_labels'] = kwargs.get('fz_labels', 14)
        kwargs['fz_cblabel'] = kwargs.get('fz_cblabel', 14)
        kwargs['fz_title'] = kwargs.get('fz_title', 16)
        sl_times = slice(edges, len(self._times) - edges)
        times = self._times[sl_times]
        pha_n = self.pha_a[..., sl_times].squeeze()
        # z-score normalization
        if zscore:
            if baseline is None:
                bsl_idx = sl_times
            else:
                assert len(baseline) == 2
                bsl_idx = slice(baseline[0], baseline[1])
            _mean = self.amp_a[..., bsl_idx].mean(2, keepdims=True)
            _std = self.amp_a[..., bsl_idx].std(2, keepdims=True)
            _std[_std == 0.] = 1.  # correction from NaN
            amp_n = (self.amp_a[..., sl_times] - _mean) / _std
        else:
            amp_n = self.amp_a[..., sl_times]

        # grid definition
        gs = GridSpec(8, 8)
        # image plot
        plt.subplot(gs[slice(0, 6), 0:-1])
        self.pacplot(amp_n.mean(1), times, self.yvec, **kwargs)
        plt.axvline(times[self.cue], color='w', lw=2)
        plt.tick_params(bottom=False, labelbottom=False)
        ax_1 = plt.gca()
        # external colorbar
        plt.subplot(gs[slice(1, 5), -1])
        cb = plt.colorbar(self._plt_im, pad=0.01, cax=plt.gca())
        cb.set_label('Power (V**2/Hz)', fontsize=kwargs['fz_cblabel'])
        cb.outline.set_visible(False)
        # phase plot
        plt.subplot(gs[slice(6, 8), 0:-1])
        plt.plot(times, pha_n.T, color='lightgray', alpha=.2, lw=1.)
        plt.plot(times, pha_n.mean(0), label='single trial phases', alpha=.2,
                 lw=1.)  # legend tweaking
        plt.plot(times, pha_n.mean(0), label='mean phases',
                 color='#1f77b4')
        plt.axvline(times[self.cue], color='k', lw=2)
        plt.autoscale(axis='both', tight=True, enable=True)
        plt.xlabel("Times", fontsize=kwargs['fz_labels'])
        plt.ylabel("V / Hz", fontsize=kwargs['fz_labels'])
        # bottom legend
        plt.legend(loc='center', bbox_to_anchor=(.5, -.5),
                   fontsize='x-large', ncol=2)
        ax_2 = plt.gca()

        return [ax_1, ax_2]