__all__ = ['BasicResidual1x', 'BasicResidual_downup_1x', 'BasicResidual2x', 'BasicResidual_downup_2x', 'FactorizedReduce', 'OPS', 'OPS_name', 'OPS_Class', 'Self_Attn', 'SAHead']

from pdb import set_trace as bp
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from thop import profile
import sys
import os.path as osp
from easydict import EasyDict as edict
from torch import nn, einsum
from einops import rearrange
import fvcore.nn.weight_init as weight_init
from detectron2.layers import Conv2d, get_norm

def drop_path_(x, drop_prob, training):
    if training and drop_prob > 0.:
        keep_prob = 1. - drop_prob
        # per data point mask; assuming x in cuda.
        mask = torch.cuda.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob)
        x = torch.div(x, keep_prob)
        x = torch.mul(x, mask)
        # x.div_(keep_prob).mul_(mask)
    return x

class DropPath_(nn.Module):
    def __init__(self, p=0.):
        """ [!] DropPath is inplace module
        Args:
            p: probability of an path to be zeroed.
        """
        super().__init__()
        self.p = p

    def extra_repr(self):
        return 'p={}, inplace'.format(self.p)

    def forward(self, x):
        drop_path_(x, self.p, self.training)

        return x

    def forward_latency(self, size):
        c_in, h_in, w_in = size
        latency = 0
        return latency, (c_in, h_in, w_in)

class BasicResidual1x(nn.Module):
    def __init__(self, norm, C_in, C_out, kernel_size=3, stride=1, dilation=1, groups=1, use_bias=False):
        super(BasicResidual1x, self).__init__()
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        groups = 1
        self.C_in = C_in
        self.C_out = C_out
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation
        self.groups = groups
        assert stride in [1, 2]
        if self.stride == 2: self.dilation = 1

        self.conv = Conv2d(
                C_in,
                C_out,
                kernel_size=3,
                stride=stride,
                padding=dilation,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out),
                activation=F.relu,
        )
        weight_init.c2_xavier_fill(self.conv)

    def forward(self, x):
        out = self.conv(x)
        return out

class BasicResidual_downup_1x(nn.Module):
    def __init__(self, norm, C_in, C_out, kernel_size=3, stride=1, dilation=1, groups=1, use_bias=False, align_corners=False):
        super(BasicResidual_downup_1x, self).__init__()
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        groups = 1
        self.C_in = C_in
        self.C_out = C_out
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation
        self.groups = groups
        assert stride in [1, 2]
        if self.stride == 2: self.dilation = 1
        self.align_corners = align_corners

        self.relu = F.relu
        self.conv = Conv2d(
                C_in,
                C_out,
                kernel_size=3,
                stride=1,
                padding=dilation,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out)
        )
        if self.stride==1:
            self.downsample = Conv2d(
                    C_in,
                    C_out,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    dilation=dilation,
                    groups=groups,
                    bias=use_bias,
                    norm=get_norm(norm, C_out)
            )
            weight_init.c2_xavier_fill(self.downsample)

        weight_init.c2_xavier_fill(self.conv)

    def forward(self, x):
        out = F.interpolate(x, size=(int(x.size(2))//2, int(x.size(3))//2), mode='bilinear', align_corners=self.align_corners)
        out = self.conv(out)
        if self.stride == 1:
            out = F.interpolate(out, size=(int(x.size(2)), int(x.size(3))), mode='bilinear', align_corners=self.align_corners)
            out = out + self.downsample(x)
        out = self.relu(out)
        return out

class BasicResidual2x(nn.Module):
    def __init__(self, norm, C_in, C_out, kernel_size=3, stride=1, dilation=1, groups=1, use_bias=False):
        super(BasicResidual2x, self).__init__()
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        groups = 1
        self.C_in = C_in
        self.C_out = C_out
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation
        self.groups = groups
        assert stride in [1, 2]
        if self.stride == 2: self.dilation = 1

        self.conv1 = Conv2d(
                C_in,
                C_out,
                kernel_size=3,
                stride=stride,
                padding=dilation,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out),
                activation=F.relu,
        )

        self.conv2 = Conv2d(
                C_out,
                C_out,
                kernel_size=3,
                stride=1,
                padding=dilation,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out),
                activation=F.relu,
        )

        weight_init.c2_xavier_fill(self.conv1)
        weight_init.c2_xavier_fill(self.conv2)

    def forward(self, x):
        x = self.conv1(x)
        out = self.conv2(x)
        return out

class BasicResidual_downup_2x(nn.Module):
    def __init__(self, norm, C_in, C_out, kernel_size=3, stride=1, dilation=1, groups=1, use_bias=False, align_corners=False):
        super(BasicResidual_downup_2x, self).__init__()
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        groups = 1
        self.C_in = C_in
        self.C_out = C_out
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation
        self.groups = groups
        self.align_corners = align_corners
        assert stride in [1, 2]
        if self.stride == 2: self.dilation = 1

        self.relu = F.relu
        self.conv1 = Conv2d(
                C_in,
                C_out,
                kernel_size=3,
                stride=1,
                padding=dilation,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out),
                activation=F.relu,
        )

        self.conv2 = Conv2d(
                C_out,
                C_out,
                kernel_size=3,
                stride=1,
                padding=dilation,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out)
        )
        weight_init.c2_xavier_fill(self.conv1)
        weight_init.c2_xavier_fill(self.conv2)

        if self.stride==1:
            self.downsample = Conv2d(
                    C_in,
                    C_out,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    dilation=dilation,
                    groups=groups,
                    bias=use_bias,
                    norm=get_norm(norm, C_out)
            )
            weight_init.c2_xavier_fill(self.downsample)

    def forward(self, x):
        out = F.interpolate(x, size=(int(x.size(2))//2, int(x.size(3))//2), mode='bilinear', align_corners=self.align_corners)
        out = self.conv1(out)
        out = self.conv2(out)
        if self.stride == 1:
            out = F.interpolate(out, size=(int(x.size(2)), int(x.size(3))), mode='bilinear', align_corners=self.align_corners)
            out = out + self.downsample(x)
        out = self.relu(out)
        return out

class FactorizedReduce(nn.Module):
    def __init__(self, norm, C_in, C_out, kernel_size=3, stride=1, dilation=1, groups=1, use_bias=False):
        super(FactorizedReduce, self).__init__()
        assert stride in [1, 2]
        assert C_out % 2 == 0
        self.C_in = C_in
        self.C_out = C_out
        self.stride = stride
        if stride == 1:
            self.conv1 = Conv2d(
                C_in,
                C_out,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out),
                activation=F.relu,
            )
            weight_init.c2_xavier_fill(self.conv1)

        elif stride == 2:
            self.conv1 = Conv2d(
                C_in,
                C_out // 2,
                kernel_size=1,
                stride=2,
                padding=0,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out // 2),
                activation=F.relu,
            )

            self.conv2 = Conv2d(
                C_in,
                C_out // 2,
                kernel_size=1,
                stride=2,
                padding=0,
                dilation=dilation,
                groups=groups,
                bias=use_bias,
                norm=get_norm(norm, C_out // 2),
                activation=F.relu,
            )
            weight_init.c2_xavier_fill(self.conv1)
            weight_init.c2_xavier_fill(self.conv2)

    def forward(self, x):
        if self.stride == 2:
            out = torch.cat([self.conv1(x), self.conv2(x[:,:,1:,1:])], dim=1)
        else:
            out = self.conv1(x)
        return out

def pair(x):
    return (x, x) if not isinstance(x, tuple) else x

def expand_dim(t, dim, k):
    t = t.unsqueeze(dim = dim)
    expand_shape = [-1] * len(t.shape)
    expand_shape[dim] = k
    return t.expand(*expand_shape)

def rel_to_abs(x):
    b, h, l, _, device, dtype = *x.shape, x.device, x.dtype
    dd = {'device': device, 'dtype': dtype}
    col_pad = torch.zeros((b, h, l, 1), **dd)
    x = torch.cat((x, col_pad), dim = 3)
    flat_x = rearrange(x, 'b h l c -> b h (l c)')
    flat_pad = torch.zeros((b, h, l - 1), **dd)
    flat_x_padded = torch.cat((flat_x, flat_pad), dim = 2)
    final_x = flat_x_padded.reshape(b, h, l + 1, 2 * l - 1)
    final_x = final_x[:, :, :l, (l-1):]
    return final_x

def relative_logits_1d(q, rel_k):
    b, heads, h, w, dim = q.shape
    logits = einsum('b h x y d, r d -> b h x y r', q, rel_k)
    logits = rearrange(logits, 'b h x y r -> b (h x) y r')
    logits = rel_to_abs(logits)
    logits = logits.reshape(b, heads, h, w, w)
    logits = expand_dim(logits, dim = 3, k = h)
    return logits

# positional embeddings

class AbsPosEmb(nn.Module):
    def __init__(
        self,
        fmap_size,
        dim_head
    ):
        super().__init__()
        height, width = pair(fmap_size)
        scale = dim_head ** -0.5
        self.height = nn.Parameter(torch.randn(height, dim_head) * scale)
        self.width = nn.Parameter(torch.randn(width, dim_head) * scale)

    def forward(self, q):
        emb = rearrange(self.height, 'h d -> h () d') + rearrange(self.width, 'w d -> () w d')
        emb = rearrange(emb, ' h w d -> (h w) d')
        logits = einsum('b h i d, j d -> b h i j', q, emb)
        return logits

class RelPosEmb(nn.Module):
    def __init__(
        self,
        fmap_size,
        dim_head
    ):
        super().__init__()
        height, width = pair(fmap_size)
        scale = dim_head ** -0.5
        self.fmap_size = fmap_size
        self.rel_height = nn.Parameter(torch.randn(height * 2 - 1, dim_head) * scale)
        self.rel_width = nn.Parameter(torch.randn(width * 2 - 1, dim_head) * scale)

    def forward(self, q):
        h, w = self.fmap_size

        q = rearrange(q, 'b h (x y) d -> b h x y d', x = h, y = w)
        rel_logits_w = relative_logits_1d(q, self.rel_width)
        rel_logits_w = rearrange(rel_logits_w, 'b h x i y j-> b h (x y) (i j)')

        q = rearrange(q, 'b h x y d -> b h y x d')
        rel_logits_h = relative_logits_1d(q, self.rel_height)
        rel_logits_h = rearrange(rel_logits_h, 'b h x i y j -> b h (y x) (j i)')
        return rel_logits_w + rel_logits_h

# classes

class Attention(nn.Module):
    def __init__(
        self,
        *,
        dim,
        fmap_size,
        heads = 4,
        dim_head = 128,
        rel_pos_emb = False
    ):
        super().__init__()
        self.heads = heads

        self.scale = dim_head ** -0.5
        inner_dim = heads * dim_head

        self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias = False)

        rel_pos_class = AbsPosEmb if not rel_pos_emb else RelPosEmb
        self.pos_emb = rel_pos_class(fmap_size, dim_head)

    def forward(self, fmap):
        heads, b, c, h, w = self.heads, *fmap.shape

        q, k, v = self.to_qkv(fmap).chunk(3, dim = 1)
        q, k, v = map(lambda t: rearrange(t, 'b (h d) x y -> b h (x y) d', h = heads), (q, k, v))

        q *= self.scale

        sim = einsum('b h i d, b h j d -> b h i j', q, k)
        sim += self.pos_emb(q)

        attn = sim.softmax(dim = -1)

        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
        return out

class Self_Attn(nn.Module):
    def __init__(
        self,
        norm,
        dim,
        fmap_size,
        dim_out,
        downsample,
        rel_pos_emb = False
    ):
        super().__init__()

        # shortcut
        if dim != dim_out or downsample:
            self.sk = False
            kernel_size, stride, padding = (3, 2, 1) if downsample else (1, 1, 0)
            self.shortcut = Conv2d(
                dim,
                dim_out,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=False,
                norm=get_norm(norm, dim_out),
                activation=F.relu,
            )
        else:
            self.sk = True
            self.shortcut = nn.Identity()

        # attn_dim_in = dim_out // proj_factor
        attn_dim_in = dim_out
        # attn_dim_out = heads * dim_head
        attn_dim_out = attn_dim_in
        activation = F.relu

        self.net1 = Conv2d(
                dim,
                attn_dim_in,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False,
                norm=get_norm(norm, attn_dim_in),
                activation=F.relu,
        )

        self.net2 = nn.Sequential(
            ATT(attn_dim_in),
            nn.AvgPool2d((2, 2)) if downsample else nn.Identity(),
            get_norm(norm, attn_dim_in),
            nn.ReLU(),
        )

        self.net3 = Conv2d(
                attn_dim_out,
                dim_out,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False,
                norm=get_norm(norm, dim_out)
        )

        weight_init.c2_xavier_fill(self.net1)
        weight_init.c2_xavier_fill(self.net3)
        # init last batch norm gamma to zero

        nn.init.zeros_(self.net3.norm.weight)

        # final activation

        self.activation = activation

    def forward(self, x):
        shortcut = self.shortcut(x)
        x = self.net1(x)
        x = self.net2(x)
        x = self.net3(x)
        x += shortcut
        return self.activation(x)

class SAHead(nn.Module):
    def __init__(
        self,
        norm,
        dim,
        fmap_size,
        dim_out,
        downsample,
        rel_pos_emb = False,
        align_corners = False
    ):
        super().__init__()

        self.stride = 2 if downsample else 1
        # shortcut
        self.align_corners = align_corners
        if dim != dim_out or downsample:
            self.sk = False
            kernel_size, stride, padding = (3, 2, 1) if downsample else (1, 1, 0)
            self.shortcut = Conv2d(
                dim,
                dim_out,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=False,
                norm=get_norm(norm, dim_out),
                activation=F.relu,
            )
        else:
            self.sk = True
            self.shortcut = nn.Identity()

        # attn_dim_in = dim_out // proj_factor
        attn_dim_in = dim_out
        # attn_dim_out = heads * dim_head
        attn_dim_out = attn_dim_in
        activation = F.relu

        self.net1 = Conv2d(
                dim,
                attn_dim_in,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False,
                norm=get_norm(norm, attn_dim_in),
                activation=F.relu,
        )

        self.net2 = nn.Sequential(
            ATT(attn_dim_in),
            nn.AvgPool2d((2, 2)) if downsample else nn.Identity(),
            get_norm(norm, attn_dim_in),
            nn.ReLU(),
        )

        self.net3 = Conv2d(
                attn_dim_out,
                dim_out,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False,
                norm=get_norm(norm, dim_out)
        )

        weight_init.c2_xavier_fill(self.net1)
        weight_init.c2_xavier_fill(self.net3)
        # init last batch norm gamma to zero

        nn.init.zeros_(self.net3.norm.weight)

        # final activation

        self.activation = activation

    def forward(self, x):
        shortcut = self.shortcut(x)
        x = self.net1(x)
        x = F.interpolate(x, size=(int(x.size(2))//2, int(x.size(3))//2), mode='bilinear', align_corners=self.align_corners)
        x = self.net2(x)
        if self.stride == 1:
            x = F.interpolate(x, size=(int(x.size(2)*2), int(x.size(3)*2)), mode='bilinear', align_corners=self.align_corners)

        x = self.net3(x)
        x += shortcut
        return self.activation(x)

class ATT(nn.Module):
    """ Self attention Layer"""
    def __init__(self, in_dim):
        super(ATT, self).__init__()
        self.chanel_in = in_dim

        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
        self.gamma = nn.Parameter(torch.zeros(1))

        self.softmax  = nn.Softmax(dim=-1) #

    def forward(self,x):
        """
            inputs :
                x : input feature maps( B X C X W X H)
            returns :
                out : self attention value + input feature
                attention: B X N X N (N is Width*Height)
        """
        m_batchsize, C, width, height = x.size()
        proj_query  = self.query_conv(x).view(m_batchsize, -1,width*height).permute(0,2,1) # B X CX(N)
        proj_key =  self.key_conv(x).view(m_batchsize,-1, width*height) # B X C x (*W*H)
        energy =  torch.bmm(proj_query, proj_key) # transpose check
        attention = self.softmax(energy) # BX (N) X (N)
        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N

        out = torch.bmm(proj_value, attention.permute(0,2,1) )
        out = out.view(m_batchsize, C, width,height)

        out = self.gamma*out + x
        return out

from collections import OrderedDict
OPS = {
    'skip' : lambda norm, C_in, C_out, stride, use_bias, fmap_size: FactorizedReduce(norm, C_in, C_out, stride, use_bias=use_bias),
    'conv' : lambda norm, C_in, C_out, stride, use_bias, fmap_size: BasicResidual1x(norm, C_in, C_out, kernel_size=3, stride=stride, dilation=1, use_bias=use_bias),
    'conv_downup' : lambda norm, C_in, C_out, stride, use_bias, fmap_size: BasicResidual_downup_1x(norm, C_in, C_out, kernel_size=3, stride=stride, dilation=1, use_bias=use_bias),
    'conv_2x' : lambda norm, C_in, C_out, stride, use_bias, fmap_size: BasicResidual2x(norm, C_in, C_out, kernel_size=3, stride=stride, dilation=1, use_bias=use_bias),
    'conv_2x_downup' : lambda norm, C_in, C_out, stride, use_bias, fmap_size: BasicResidual_downup_2x(norm, C_in, C_out, kernel_size=3, stride=stride, dilation=1, use_bias=use_bias),
    'sa': lambda norm, C_in, C_out, stride, use_bias, fmap_size: Self_Attn(norm, dim=C_in, fmap_size=(128, 256), dim_out=C_out, downsample=(stride==2))
}

OPS_name = ["FactorizedReduce", "BasicResidual1x", "BasicResidual_downup_1x", "BasicResidual2x", "BasicResidual_downup_2x", "Self_Attn"]

OPS_Class = OrderedDict()
OPS_Class['skip'] = FactorizedReduce
OPS_Class['conv'] = BasicResidual1x
OPS_Class['conv_downup'] = BasicResidual_downup_1x
OPS_Class['conv_2x'] = BasicResidual2x
OPS_Class['conv_2x_downup'] = BasicResidual_downup_2x
OPS_Class['sa'] = Self_Attn


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

import numpy as np

from qutip import (tensor, identity, destroy, sigmax, sigmaz, basis, Qobj,
                    QobjEvo)
from ..circuit import QubitCircuit
from ..operations import Gate
from .processor import Processor
from .modelprocessor import ModelProcessor
from ..operations import expand_operator
from ..pulse import Pulse
from ..compiler import (GateCompiler, CavityQEDCompiler)

__all__ = ['DispersiveCavityQED']

class DispersiveCavityQED(ModelProcessor):
    """
    The processor based on the physical implementation of
    a dispersive cavity QED system.
    The available Hamiltonian of the system is predefined.
    For a given pulse amplitude matrix, the processor can
    calculate the state evolution under the given control pulse,
    either analytically or numerically.
    (Only additional attributes are documented here, for others please
    refer to the parent class :class:`.ModelProcessor`)

    Parameters
    ----------
    num_qubits: int
        The number of qubits in the system.

    correct_global_phase: float, optional
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.

    num_levels: int, optional
        The number of energy levels in the resonator.

    t1: list or float, optional
        Characterize the decoherence of amplitude damping for
        each qubit. A list of size `num_qubits` or a float for all qubits.

    t2: list of float, optional
        Characterize the decoherence of dephasing for
        each qubit. A list of size `num_qubits` or a float for all qubits.

    **params:
        Keyword argument for hardware parameters, in the unit of GHz.
        Qubit parameters can either be a float or a list of the length
        ``num_qubits``.

        - ``deltamax``: the pulse strength of sigma-x control, default ``1.0``
        - ``epsmax``: the pulse strength of sigma-z control, default ``9.5``
        - ``eps``: the bare transition frequency for each of the qubits,
          default ``9.5``
        - ``delta``: the coupling between qubit states, default ``0.0``
        - ``g``: the coupling strength between the resonator and the qubit,
          default ``1.0``
        - ``w0``: the bare frequency of the resonator. Should only be a float,
          default ``0.01``

        The dressed qubit frequency is `wq` is computed by
        :math:`w_q=\sqrt{\epsilon^2+\delta^2}`

    Attributes
    ----------
    wq: list of float
        The frequency of the qubits calculated from
        eps and delta for each qubit.

    Delta: list of float
        The detuning with respect to w0 calculated
        from wq and w0 for each qubit.
    """

    def __init__(self, num_qubits, correct_global_phase=True,
                 num_levels=10, t1=None, t2=None, **params):
        super(DispersiveCavityQED, self).__init__(
            num_qubits, correct_global_phase=correct_global_phase,
            t1=t1, t2=t2)
        self.correct_global_phase = correct_global_phase
        self.spline_kind = "step_func"
        self.num_levels = num_levels
        self.params = {  # default parameters
            "deltamax": 1.0,
            "epsmax": 9.5,
            "w0": 10,
            "eps": 9.5,
            "delta": 0.0,
            "g": 0.01,
        }
        if params is not None:
            self.params.update(params)
        for key, value in self.params.items():
            if key != "w0":
                # if float, make it an array
                self.params[key] = self.to_array(value, self.num_qubits)
        self.set_up_params()
        self.set_up_ops(num_qubits)
        self.dims = [num_levels] + [2] * num_qubits
        self.pulse_dict = self.get_pulse_dict()
        self.native_gates = ["SQRTISWAP", "ISWAP", "RX", "RZ"]

    def set_up_ops(self, num_qubits):
        """
        Generate the Hamiltonians for the spinchain model and save them in the
        attribute `ctrls`.

        Parameters
        ----------
        num_qubits: int
            The number of qubits in the system.
        """
        # single qubit terms
        for m in range(num_qubits):
            self.add_control(2*np.pi*sigmax(), [m+1], label="sx" + str(m))
        for m in range(num_qubits):
            self.add_control(2*np.pi*sigmaz(), [m+1], label="sz" + str(m))
        # coupling terms
        a = tensor(
            [destroy(self.num_levels)] +
            [identity(2) for n in range(num_qubits)])
        for n in range(num_qubits):
            sm = tensor([identity(self.num_levels)] +
                        [destroy(2) if m == n else identity(2)
                         for m in range(num_qubits)])
            self.add_control(
                2*np.pi * a.dag() * sm + 2*np.pi * a * sm.dag(),
                list(range(num_qubits+1)), label="g" + str(n)
            )

    def set_up_params(self):
        """
        Compute the qubit frequency and detune.
        """
        # backward compatibility
        self.params["sz"] = self.params["epsmax"]
        self.params["sx"] = self.params["deltamax"]
        eps = self.params["eps"]
        delta = self.params["delta"]
        w0 = self.params["w0"]
        g = self.params["g"]
        # computed
        self.wq = np.sqrt(eps**2 + delta**2)
        self.Delta = self.wq - w0

        # rwa/dispersive regime tests
        if any(g / (w0 - self.wq) > 0.05):
            warnings.warn("Not in the dispersive regime")

        if any((w0 - self.wq)/(w0 + self.wq) > 0.05):
            warnings.warn(
                "The rotating-wave approximation might not be valid.")
        print(self.params)

    @property
    def sx_ops(self):
        """
        list: A list of sigmax Hamiltonians for each qubit.
        """
        return self.ctrls[0: self.num_qubits]

    @property
    def sz_ops(self):
        """
        list: A list of sigmaz Hamiltonians for each qubit.
        """
        return self.ctrls[self.num_qubits: 2*self.num_qubits]

    @property
    def cavityqubit_ops(self):
        """
        list: A list of interacting Hamiltonians between cavity and each qubit.
        """
        return self.ctrls[2*self.num_qubits: 3*self.num_qubits]

    @property
    def sx_u(self):
        """array-like: Pulse matrix for sigmax Hamiltonians."""
        return self.coeffs[: self.num_qubits]

    @property
    def sz_u(self):
        """array-like: Pulse matrix for sigmaz Hamiltonians."""
        return self.coeffs[self.num_qubits: 2*self.num_qubits]

    @property
    def g_u(self):
        """
        array-like: Pulse matrix for interacting Hamiltonians
        between cavity and each qubit.
        """
        return self.coeffs[2*self.num_qubits: 3*self.num_qubits]

    def get_operators_labels(self):
        """
        Get the labels for each Hamiltonian.
        It is used in the method method :meth:`.Processor.plot_pulses`.
        It is a 2-d nested list, in the plot,
        a different color will be used for each sublist.
        """
        return ([[r"$\sigma_x^%d$" % n for n in range(self.num_qubits)],
                 [r"$\sigma_z^%d$" % n for n in range(self.num_qubits)],
                 [r"$g_{%d}$" % (n) for n in range(self.num_qubits)]])

    def eliminate_auxillary_modes(self, U):
        """
        Eliminate the auxillary modes like the cavity modes in cqed.
        """
        psi_proj = tensor(
            [basis(self.num_levels, 0)] +
            [identity(2) for n in range(self.num_qubits)])
        return psi_proj.dag() * U * psi_proj

    def load_circuit(
            self, qc, schedule_mode="ASAP", compiler=None):
        if compiler is None:
            compiler = CavityQEDCompiler(
                self.num_qubits, self.params, global_phase=0.)
        tlist, coeff = super().load_circuit(
            qc, schedule_mode=schedule_mode, compiler=compiler)
        self.global_phase = compiler.global_phase
        return tlist, coeff


"""
Base class for the generative / discriminative models.

Copyright (C) 2017, Lucas Ondel

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""

import abc
import numpy as np
from ..io.persistent_model import PersistentModel

class EFDStats(object):
    """Wrapper of a list of accumulated sufficient statistics."""

    def __init__(self, stats):
        """Initialize the list of stats.

        Parameters
        ----------
        stats : list
            List of numpy.ndarray.

        """
        self._stats = stats

    def __setitem__(self, k, value):
        self._stats[k] = value

    def __getitem__(self, k):
        return self._stats[k]

    def __iadd__(self, stats):
        for i in range(len(self._stats)):
            self._stats[i] += stats[i]
        return self

    def __imul__(self, scale):
        for i in range(len(self._stats)):
            self._stats[i] *= scale
        return self

class DiscreteLatentModel(PersistentModel, metaclass=abc.ABCMeta):
    """Abstract base class models with latent discrete variable.

    Concrete implementation of this interface can be either
    a Mixture Model (when the latent variables are independent) or
    a Hidden Markov Model (when the latent variables follow a 1st
    order Markov dynamic).

    NOTE
    ----
    This implementation assume the joint distribution of the
    data and the latent variables to be a member of the Exponential
    family of distribution.

    """

    def __init__(self, latent_prior, latent_posterior, components):
        """Initialize a :class:`DiscreteLatentModel`

        Parameters
        ----------
        latent_prior : :class:`EFDPrior`
            Prior over the latent variable.
        latent_posterior : :class:`EFDPrior`
            Initial guess for the posterior over the latent variable.
        components : list or tuple of :class:`EFDLikelihood`
            List of distributions / densities for the conditional
            likelihood. No check is performed but we assumed that all
            the elements of the list are of the same distribution /
            density type.

        """
        self._latent_prior = latent_prior
        self._latent_posterior = latent_posterior
        self._components = components

        # Pre-compute the expected value of the natural parameters for
        # each components to speed computation during
        # training/decoding.
        self._exp_np_matrix = self._get_components_params_matrix()

    @property
    def latent_prior(self):
        return self._latent_prior

    @latent_prior.setter
    def latent_prior(self, value):
        self._latent_prior = value

    @property
    def latent_posterior(self):
        return self._latent_posterior

    @latent_posterior.setter
    def latent_posterior(self, value):
        self._latent_posterior = value

    @property
    def components(self):
        return self._components

    @components.setter
    def components(self, value):
        self._components = value

    def _get_components_params_matrix(self):
        """Natural parameters of the components as a matrix.

        Expected value of the natural parameters w.r.t. their current
        posterior distribution organized as a matrix.

        Returns
        -------
        exp_np_matrix :
            Expected value of the natural parameters of the components
             as a matrix.

        """
        return np.vstack([comp.posterior.grad_log_partition
                          for comp in self.components])

    def get_sufficient_stats(self, data):
        """Sufficient statistics of the latent model.

        Parameters
        ----------
        data : numpy.ndarray
            (N x D) matrix where N is the number of frames and D is the
            dimension of a single features vector.

        Returns
        -------
        s_stats : numpy.ndarray
            (N x D2) matrix of sufficient statistics. D2 is the
            dimension of the sufficient statistics for a single
            features frame.

        """
        cls = self.components[0].__class__
        s_stats = cls.get_sufficient_stats(data)
        return s_stats

    def components_exp_llh(self, s_stats):
        """Per components expected log-likelihood.

        Parameters
        ----------
        s_stats : numpy.ndarray
            (NxD) matrix of sufficient statistics.

        """
        return self._exp_np_matrix.dot(s_stats.T)

    def natural_grad_update(self, acc_stats, lrate):
        """Natural gradient update.

        Parameters
        ----------
        acc_stats : :class:`EFDStats`
            Accumulated sufficient statistics.
        lrate : float
            Learning rate.

        """
        grad = self.latent_prior.natural_params + acc_stats[0]
        grad -= self.latent_posterior.natural_params
        self.latent_posterior.natural_params = \
            self.latent_posterior.natural_params + lrate * grad

        for idx, stats in enumerate(acc_stats[1]):
            comp = self.components[idx]
            grad = comp.prior.natural_params + stats
            grad -= comp.posterior.natural_params
            comp.posterior.natural_params = \
                comp.posterior.natural_params + lrate * grad

        self.post_update()

    def kl_div_posterior_prior(self):
        """Sum of KL divergence between posteriors/priors.

        The sum of KL divergence between posterior/prior of the
        parameters of the model.

        Returns
        -------
        kl_div : float
            Kullback-Leibler divergence.
        """
        retval = 0.
        retval += self.latent_posterior.kl_div(self.latent_prior)
        for comp in self.components:
            retval += comp.posterior.kl_div(comp.prior)
        return retval

    def post_update(self):
        """Called after each update of the parameters."""
        self._exp_np_matrix = self._get_components_params_matrix()

    @abc.abstractmethod
    def get_posteriors(self, s_stats, accumulate=False):
        """Compute the posterior distribution of the latent variables.

        Parameters
        ----------
        s_stats : numpy.ndarray
            (N x D) matrix of sufficient statistics.
        accumulate : boolean
            If True, accumulate the sufficient statistics based on the
            posteriors.

        Returns
        -------
        posts : numpy.ndarray
            (N x K) posterior matrix. K is the number of state of the
            latent variables.
        log_evidence : float
            Log evidence of the data. For most, if not all, of the
            models the evidence is intractable and a lower-bound is
            returned instead.
        acc_stats : :class:`EFDStats`
            If accumulate is True, return accumulated sufficent
            statistics.

        """
        pass

    # PersistentModel interface implementation.
    # -----------------------------------------------------------------

    def to_dict(self):
        return {
            'latent_prior_class': self.latent_prior.__class__,
            'latent_prior_data': self.latent_prior.to_dict(),
            'posterior_class': self.latent_posterior.__class__,
            'posterior_data': self.latent_posterior.to_dict(),
            'components_class': self.components[0].__class__,
            'components': [comp.to_dict() for comp in components]
        }

    @staticmethod
    def load_from_dict(cls, model_data):
        model = cls.__new__(model_data['class'])
        latent_prior_cls = model_data['latent_prior_class']
        latent_prior_data = model_data['latent_prior_data']
        model.latent_prior = \
            latent_prior_cls.load_from_dict(latent_prior_data)

        latent_posterior_cls = model_data['latent_posterior_class']
        latent_posterior_data = model_data['latent_posterior_data']
        model.latent_posterior = \
            latent_posterior_cls.load_from_dict(latent_posterior_data)

        components = []
        components_class = model_data['components_class']
        for comp_data in model_data['components']:
            comp = components_class.load_from_dict(comp_data)
        model.components = components

        model.post_update()

        return model

    # -----------------------------------------------------------------


from abc import abstractmethod

import dask
import dask.dataframe as dd
import hiurlparser
import numpy as np
from dask.delayed import delayed
from dask_ml.preprocessing import MinMaxScaler, StandardScaler
from sklearn.preprocessing import MaxAbsScaler
from dask_ml.impute import SimpleImputer
from optimus.engines.base.distributed.functions import DistributedBaseFunctions
from optimus.helpers.core import one_tuple_to_val

class DaskBaseFunctions(DistributedBaseFunctions):
    _engine = dask

    @staticmethod
    @property
    def _functions(self):
        return dd

    @staticmethod
    def _new_series(series, *args, **kwargs):
        if isinstance(series, dd.Series):
            return series
        return dd.from_array(series, *args, **kwargs)

    @staticmethod
    def compute(*args, **kwargs):
        result = dask.compute(*(*(a for a in args), *(kwargs[k] for k in kwargs)))
        return one_tuple_to_val(result)

    @abstractmethod
    def from_dataframe(self, dfd):
        pass

    @staticmethod
    def to_dataframe(dfd):
        return dfd.compute()

    @staticmethod
    def df_concat(df_list):
        return dd.concat(df_list, axis=0)

    def new_df(self, *args, **kwargs):
        if len(args) < 4 and all([k not in kwargs for k in ['dsk', 'name', 'meta', 'divisions']]):
            return self.from_dataframe(self._partition_engine.DataFrame(*args, **kwargs))
        return dd.DataFrame(*args, **kwargs)

    @staticmethod
    def dask_to_compatible(dfd):
        return dfd

    def sort_df(self, dfd, cols, ascending):
        for c, a in list(zip(cols, ascending))[::-1]:
            dfd = dfd.sort_values(c)
            if not a:
                dfd = self.reverse_df(dfd)
        return dfd.reset_index(drop=True)

    def reverse_df(self, dfd):
        @delayed
        def reverse_pdf(pdf):
            return pdf[::-1]

        ds = dfd.to_delayed()
        ds = [reverse_pdf(d) for d in ds][::-1]
        return dd.from_delayed(ds)

    def to_float(self, series):
        if getattr(series, "map_partitions", False):
            return self.map_partitions(series, self._to_float)
        else:
            return self._to_float(series)

    def to_integer(self, series, default=0):
        if getattr(series, "map_partitions", False):
            return self.map_partitions(series, self._to_integer, default=default)
        else:
            return self._to_integer(series, default=default)

    def to_datetime(self, series):
        if getattr(series, "map_partitions", False):
            return self.map_partitions(series, self._to_datetime)
        else:
            return self._to_float(series)

    def duplicated(self, dfd, keep, subset):
        return self.from_dataframe(self.to_dataframe(dfd).duplicated(keep=keep, subset=subset))

    def impute(self, series, strategy, fill_value):

        imputer = SimpleImputer(strategy=strategy, fill_value=fill_value)
        series_fit = series.dropna()
        if str(series.dtype) in self.constants.OBJECT_TYPES:
            series_fit = series_fit.astype(str)
        values = series_fit.values.reshape(-1, 1)
        if len(values):
            imputer.fit(values)
            return imputer.transform(series.fillna(np.nan).values.reshape(-1, 1))
        else:
            logger.warn("list to fit imputer is empty, try cols.fill_na instead.")
            return series

    @staticmethod
    def delayed(func):
        def wrapper(*args, **kwargs):
            return dask.delayed(func)(*args, **kwargs)

        return wrapper

    def from_delayed(self, delayed):
        return dd.from_delayed(delayed)

    def to_delayed(self, value):
        return value.to_delayed()

    def mode(self, series):
        # Uses delayed wrapper to avoid unexpected Dask errors
        # TODO should this be done in every static function called from Dask instead?
        @self.delayed
        def compute_mode(series):
            # dfd = series.rename(0).value_counts().reset_index(drop=False)
            # _max = dfd[0].max(skipna=True)
            # return dfd[dfd[0] == _max]['index'].rename(series.name)
            return series.mode()

        return compute_mode(series)

    def count_zeros(self, series):
        return int((self.to_float(series).values == 0).sum())

    def standard_scaler(self, series):
        dfd = StandardScaler().fit_transform(self.to_float(series).to_frame())
        return dfd[dfd.columns[0]]

    def max_abs_scaler(self, series):
        return MaxAbsScaler().fit_transform(self.compute(self.to_float(series)).values.reshape(-1, 1))

    def min_max_scaler(self, series):
        dfd = MinMaxScaler().fit_transform(self.to_float(series).to_frame())
        return dfd[dfd.columns[0]]

    # @staticmethod
    # def heatmap(df, bins):
    #     counts, edges = da.histogramdd((df['x'], df['y'].values), bins=bins)
    #     return counts, edges[0], edges[1]

    def domain(self, series):

        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["domain"], na_action=None,
                                          meta=(series.name, "str"))

    def top_domain(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["top_domain"], na_action=None,
                                          meta=(series.name, "str"))

    def sub_domain(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["sub_domain"], na_action=None,
                                          meta=(series.name, "str"))

    def url_scheme(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["protocol"], na_action=None,
                                          meta=(series.name, "str"))

    def url_path(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["path"], na_action=None,
                                          meta=(series.name, "str"))

    def url_file(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["file"], na_action=None,
                                          meta=(series.name, "str"))

    def url_query(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["query"], na_action=None,
                                          meta=(series.name, "str"))

    def url_fragment(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["fragment"], na_action=None,
                                          meta=(series.name, "str"))

    def host(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["host"], na_action=None,
                                          meta=(series.name, "str"))

    def port(self, series):
        return self.to_string(series).map(lambda v: hiurlparser.parse_url(v)["port"], na_action=None,
                                          meta=(series.name, "str"))


import warnings

import astropy.units as u
import numpy as np
from astropy.time import Time

from sora.config.decorators import deprecated_function, deprecated_alias
from sora.prediction import occ_params, PredictionTable
from .fitting import fit_ellipse as ellipse_fitting

__all__ = ['Occultation']
warnings.simplefilter('always', UserWarning)

class Occultation:
    """Instantiates the Occultation object and performs the reduction of the
    occultation.

    Attributes
    ----------
    star : `sora.Star`, `str`, required
        the coordinate of the star in the same reference frame as the ephemeris.
        It must be a Star object or a string with the coordinates of the object
        to search on Vizier.

    body : `sora.Body`, `str`
        Object that will occult the star. It must be a Body object or its name
        to search in the Small Body Database.

    ephem : `sora.Ephem`, `list`
        Object ephemeris. It must be an Ephemeris object or a list.

    time : `str`, `astropy.time.Time`, required
        Reference time of the occultation. Time does not need to be exact, but
        needs to be within approximately 50 minutes of the occultation closest
        approach to calculate occultation parameters.

    reference_center : `str`, `sora.Observer`, `sora.Spacecraft`
        A SORA observer object or a string 'geocenter'.
        The occultation parameters will be calculated in respect
        to this reference as center of projection.

    Important
    ---------
    When instantiating with "body" and "ephem", the user may define the
    Occultation in 3 ways:

    1. With `body` and `ephem`.

    2. With only "body". In this case, the "body" parameter must be a Body
    object and have an ephemeris associated (see Body documentation).

    3. With only `ephem`. In this case, the `ephem` parameter must be one of the
    Ephem Classes and have a name (see Ephem documentation) to search for the
    body in the Small Body Database.

    """
    def __init__(self, star, body=None, ephem=None, time=None, reference_center='geocenter'):

        from sora.body import Body
        from sora.star import Star
        from .chordlist import ChordList

        if body is None and ephem is None:
            raise ValueError('"body" and/or "ephem" must be given.')
        if time is None:
            raise ValueError('"time" parameter must be given.')
        if isinstance(star, str):
            star = Star(coord=star)
        elif not isinstance(star, Star):
            raise ValueError('"star" must be a Star object or a string with coordinates of the star')
        self._star = star
        if body is not None:
            if not isinstance(body, (str, Body)):
                raise ValueError('"body" must be a string with the name of the object or a Body object')
            if isinstance(body, str):
                body = Body(name=body)
            self._body = body
        if ephem is not None:
            if body is not None:
                self.body.ephem = ephem
            else:
                if hasattr(ephem, 'name'):
                    self._body = Body(name=ephem.name, ephem=ephem)
                else:
                    raise ValueError('When only "ephem" is given, "ephem" must have a name for search.')
        try:
            ephem = self.body.ephem
        except AttributeError:
            raise ValueError('An Ephem object must be defined in Body object.')
        self._reference_center = reference_center

        tca, ca, pa, vel, dist = occ_params(self.star, ephem, time, reference_center=reference_center)
        self.ca = ca   # Closest Approach distance
        self.pa = pa   # Position Angle at CA
        self.vel = vel  # Shadow velocity at CA
        self.dist = dist  # object distance at CA
        self.tca = tca   # Instant of CA
        self.star_diam = self.star.apparent_diameter(self.dist, verbose=False)

        meta = {
            'name': self.body.name, 'radius': self.body.radius.to(u.km).value,
            'error_ra': self.body.ephem.error_ra.to(u.mas).value, 'error_dec': self.body.ephem.error_dec.to(u.mas).value}
        self.predict = PredictionTable(
            time=[tca], coord_star=[self.star.get_position(tca, observer=reference_center)],
            coord_obj=[self.body.ephem.get_position(tca, observer=reference_center)],
            ca=[ca.value], pa=[pa.value], vel=[vel.value], dist=[dist.value],
            mag=[self.star.mag['G']], source=[self.star.code], meta=meta)

        self.__observations = []
        self._chords = ChordList(star=self.star, body=self._body, time=self.tca)
        self._chords._shared_with['occultation'] = {"vel": np.absolute(self.vel), "dist": float(self.dist.AU),
                                                    "star_diam": float(self.star_diam.km)}

    @property
    def star(self):
        return self._star

    @property
    def body(self):
        return self._body

    @property
    def chords(self):
        return self._chords

    # remove this block for v1.0
    @deprecated_function(message="Please use chords.add_chord to add new observations.")
    def add_observation(self, obs, lightcurve):
        """Adds observations to the Occultation object.

        Parameters
        ----------
        obs : `sora.Observer`
            The Observer object to be added.

        lightcurve : `sora.LightCurve`
            The LightCurve object to be added.
        """
        self.chords.add_chord(name=lightcurve.name, observer=obs, lightcurve=lightcurve)

    @deprecated_function(message="Please use chords.remove_chord to remove observations.")
    def remove_observation(self, key, key_lc=None):
        """Removes an observation from the Occultation object.

        Parameters
        ----------
        key : `str`
            The name given to `Observer` or `LightCurve` to remove from the list.

        key_lc : `str`
            In the case where repeated names are present for different
            observations, `key_lc` must be given for the name of the `LightCurve`
            and key will be used for the name of the `Observer`.
        """
        try:
            self.chords.remove_chord(name=key)
        except KeyError:
            rm_list = []
            for name, chord in self.chords.items():
                if chord.observer.name == key and (key_lc is None or chord.lightcurve.name == key_lc):
                    rm_list.append(name)
            if len(rm_list) != 1:
                if len(rm_list) == 0:
                    err_mes = "No observation was identified with given keys."
                else:
                    err_mes = "More than one chord was identified. Please provide unique keys."
                raise ValueError(err_mes)
            else:
                self.chords.remove_chord(name=rm_list[0])

    @deprecated_function(message="Please use chords.")
    def observations(self):
        """ Print all the observations added to the Occultation object
        Pair (`Observer`, `LightCurve`)
        """
        print(self.chords.__repr__())

    # end of block removal

    def fit_ellipse(self, **kwargs):
        chisquare = ellipse_fitting(self, **kwargs)
        return chisquare

    fit_ellipse.__doc__ = ellipse_fitting.__doc__

    # remove this block for v1.0
    @property
    @deprecated_function(message="Please use chords.summary()")
    def positions(self):
        """ Calculates the position and velocity for all chords.
        Saves it into an `_PositionDict` object.
        """
        from functools import partial
        from .meta import _PositionDict

        if not hasattr(self, '_position'):
            self._position = _PositionDict()
        position = self._position
        if len(self.chords) == 0:
            raise ValueError('There is no observation defined for this occultation')

        pair = []
        for name, chord in self.chords.items():
            obs = chord.observer
            obs_name = obs.name
            if obs_name == '':
                obs_name = '{}_obs'.format(chord.name)
            lc = chord.lightcurve
            lc_name = lc.name
            if lc_name == '':
                lc_name = '{}_lc'.format(chord.name)
            pair.append((obs_name, lc_name))

            coord = [obs.lon, obs.lat, obs.height]
            if obs.name not in position.keys():
                position['_occ_'+obs_name] = _PositionDict(lon=obs.lon, lat=obs.lat, height=obs.height)
                position[obs_name]['_occ_lon'] = obs.lon
                position[obs_name]['_occ_lat'] = obs.lat
                position[obs_name]['_occ_height'] = obs.height
            pos_obs = position[obs_name]
            coord2 = [pos_obs['lon'], pos_obs['lat'], pos_obs['height']]
            if obs.lon != pos_obs['lon']:
                position[obs_name]['_occ_lon'] = obs.lon
            if obs.lat != pos_obs['lat']:
                position[obs_name]['_occ_lat'] = obs.lat
            if obs.height != pos_obs['height']:
                position[obs_name]['_occ_height'] = obs.height
            samecoord = (coord == coord2)

            if lc_name not in pos_obs.keys():
                pos_obs['_occ_'+lc_name] = _PositionDict()
            pos_lc = pos_obs[lc_name]

            pos_lc['_occ_status'] = chord.status()

            if hasattr(lc, 'immersion'):
                if 'immersion' not in pos_lc.keys():
                    pos_lc['_occ_immersion'] = _PositionDict(on=chord.is_able['immersion'],
                                                             enable=partial(chord.enable, time='immersion'),
                                                             disable=partial(chord.disable, time='immersion'))
                obs_im = pos_lc['immersion']
                obs_im['_occ_on'] = chord.is_able['immersion']
                do_err = False
                if samecoord and 'time' in obs_im.keys() and obs_im['time'] == lc.immersion:
                    pass
                else:
                    do_err = True
                    f1, g1, vf1, vg1 = chord.get_fg(time='immersion', vel=True)
                    obs_im['_occ_time'] = lc.immersion
                    obs_im['_occ_value'] = (round(f1, 3), round(g1, 3))
                    obs_im['_occ_vel'] = (round(vf1, 3), round(vg1, 3))
                if not do_err and 'time_err' in obs_im.keys() and obs_im['time_err'] == lc.immersion_err:
                    pass
                else:
                    fe1, ge1 = chord.get_fg(time=lc.immersion-lc.immersion_err*u.s)
                    fe2, ge2 = chord.get_fg(time=lc.immersion+lc.immersion_err*u.s)
                    obs_im['_occ_time_err'] = lc.immersion_err
                    obs_im['_occ_error'] = ((round(fe1, 3), round(ge1, 3)), (round(fe2, 3), round(ge2, 3)))

            if hasattr(lc, 'emersion'):
                pos_lc['_occ_emersion'] = _PositionDict(on=chord.is_able['emersion'],
                                                        enable=partial(chord.enable, time='emersion'),
                                                        disable=partial(chord.disable, time='emersion'))
                obs_em = pos_lc['emersion']
                do_err = False
                if samecoord and 'time' in obs_em.keys() and obs_em['time'] == lc.emersion:
                    pass
                else:
                    do_err = True
                    f1, g1, vf1, vg1 = chord.get_fg(time='emersion', vel=True)
                    obs_em['_occ_time'] = lc.emersion
                    obs_em['_occ_value'] = (round(f1, 3), round(g1, 3))
                    obs_em['_occ_vel'] = (round(vf1, 3), round(vg1, 3))
                if not do_err and 'time_err' in obs_em.keys() and obs_em['time_err'] == lc.emersion_err:
                    pass
                else:
                    fe1, ge1 = chord.get_fg(time=lc.emersion-lc.emersion_err*u.s)
                    fe2, ge2 = chord.get_fg(time=lc.emersion+lc.emersion_err*u.s)
                    obs_em['_occ_time_err'] = lc.emersion_err
                    obs_em['_occ_error'] = ((round(fe1, 3), round(ge1, 3)), (round(fe2, 3), round(ge2, 3)))

            if pos_lc['status'] == 'negative':
                if 'start_obs' not in pos_lc.keys():
                    pos_lc['_occ_start_obs'] = _PositionDict()
                obs_start = pos_lc['start_obs']
                if samecoord and 'time' in obs_start.keys() and obs_start['time'] == lc.initial_time:
                    pass
                else:
                    f, g, vf, vg = chord.get_fg(time='start', vel=True)
                    obs_start['_occ_time'] = lc.initial_time
                    obs_start['_occ_value'] = (round(f, 3), round(g, 3))
                    obs_start['_occ_vel'] = (round(vf, 3), round(vg, 3))
                if 'end_obs' not in pos_lc.keys():
                    pos_lc['_occ_end_obs'] = _PositionDict()
                obs_end = pos_lc['end_obs']
                if samecoord and 'time' in obs_end.keys() and obs_end['time'] == lc.end_time:
                    pass
                else:
                    f, g, vf, vg = chord.get_fg(time='end', vel=True)
                    obs_end['_occ_time'] = lc.end_time
                    obs_end['_occ_value'] = (round(f, 3), round(g, 3))
                    obs_end['_occ_vel'] = (round(vf, 3), round(vg, 3))

        for key in list(position):
            n = 0
            for key_lc in list(position[key]):
                if type(position[key][key_lc]) != _PositionDict:
                    continue
                if (key, key_lc) not in pair:
                    del position[key][key_lc]
                else:
                    n += 1
            if n == 0:
                del position[key]

        return self._position

    @positions.setter
    def positions(self, value):
        """
        Note
        ----
        If the users tries to set a value to position, it must be ``'on'`` or
        ``'off'``, and it will be assigned to all chords.
        """
        if value not in ['on', 'off']:
            raise ValueError("Value must be 'on' or 'off' only.")
        pos = self.positions
        for key in pos.keys():
            pos[key] = value

    # end of block removal

    def check_velocities(self):
        """Prints the current velocity used by the LightCurves and its radial velocity.
        """
        if hasattr(self, 'fitted_params'):
            center = np.array([self.fitted_params['center_f'][0], self.fitted_params['center_g'][0]])
        else:
            center = np.array([0, 0])
        for name, chord in self.chords.items():
            im = getattr(chord.lightcurve, 'immersion', None)
            em = getattr(chord.lightcurve, 'emersion', None)
            if im is None and em is None:
                continue
            print('{} - Velocity used: {:.3f}'.format(name, chord.lightcurve.vel))
            if im is not None:
                vals = chord.get_fg(time=im, vel=True)
                delta = np.array(vals[0:2]) - center
                print('    Immersion Radial Velocity: {:.3f}'.
                      format(np.abs(np.dot(np.array(vals[2:]), delta)/np.linalg.norm(delta))))
            if em is not None:
                vals = chord.get_fg(time=em, vel=True)
                delta = np.array(vals[0:2]) - center
                print('    Emersion Radial Velocity: {:.3f}'.
                      format(np.abs(np.dot(np.array(vals[2:]), delta)/np.linalg.norm(delta))))

    @deprecated_alias(log='verbose')  # remove this line in v1.0
    def new_astrometric_position(self, time=None, offset=None, error=None, verbose=True, observer=None):
        """Calculates the new astrometric position for the object given fitted parameters.

        Parameters
        ----------
        time : `str`, `astropy.time.Time`
            Reference time to calculate the position. If not given, it uses the
            instant of the occultation Closest Approach. It can be a string
            in the ISO format (yyyy-mm-dd hh:mm:ss.s) or an astropy Time object.

        offset : `list`
            Offset to apply to the position. If not given, uses the parameters
            from the fitted ellipse.

            Note
            ----
            Must be a list of 3 values being [X, Y, 'unit']. 'unit' must be
            Angular or Distance unit.

            If Distance units for X and Y:
            Ex: [100, -200, 'km'], [0.001, 0.002, 'AU']

            If Angular units fox X [d*a*cos(dec)] and Y [d*dec]:
            Ex: [30.6, 20, 'mas'], or [-15, 2, 'arcsec']

        error : `list`
            Error bar of the given offset. If not given, it uses the 1-sigma
            value of the fitted ellipse.

            Note
            ----
            Error must be a list of 3 values being [dX, dY, 'unit'], similar to
            offset. It does not need to be in the same unit as offset.

        verbose : `bool`
            If true, it Prints text, else it Returns text.

        observer : `str`, `sora.Observer`, `sora.Spacecraft`
            IAU code of the observer (must be present in given list of kernels),
            a SORA observer object or a string: ['geocenter', 'barycenter']
        """
        from astropy.coordinates import SkyCoord, SkyOffsetFrame

        if time is not None:
            time = Time(time)
        else:
            time = self.tca

        tca_diff = np.absolute(time-self.tca)
        if tca_diff > 1*u.day:
            warnings.warn('The difference between the given time and the closest approach instant is {:.1f} days. '
                          'This position could not have a physical meaning.'.format(tca_diff.jd))

        if offset is not None:
            off_ra = offset[0]*u.Unit(offset[2])
            off_dec = offset[1]*u.Unit(offset[2])
            if off_ra.unit.physical_type == 'length':
                dist = True
            elif off_ra.unit.physical_type == 'angle':
                dist = False
            else:
                raise ValueError('Offset unit must be a distance or angular value.')
        elif hasattr(self, 'fitted_params'):
            off_ra = self.fitted_params['center_f'][0]*u.km
            off_dec = self.fitted_params['center_g'][0]*u.km
            dist = True
        else:
            warnings.warn('No offset given or found. Using 0.0 instead.')
            off_ra = 0.0*u.mas
            off_dec = 0.0*u.mas
            dist = False

        if error is not None:
            e_off_ra = error[0]*u.Unit(error[2])
            e_off_dec = error[1]*u.Unit(error[2])
            if e_off_ra.unit.physical_type == 'length':
                e_dist = True
            elif e_off_ra.unit.physical_type == 'angle':
                e_dist = False
            else:
                raise ValueError('Error unit must be a distance or angular value.')
        elif hasattr(self, 'fitted_params'):
            e_off_ra = self.fitted_params['center_f'][1]*u.km
            e_off_dec = self.fitted_params['center_g'][1]*u.km
            e_dist = True
        else:
            warnings.warn('No error given or found. Using 0.0 instead.')
            e_off_ra = 0.0*u.mas
            e_off_dec = 0.0*u.mas
            e_dist = False

        if observer is None:
            observer = self._reference_center

        coord = self.body.ephem.get_position(time, observer=observer)
        distance = coord.distance.to(u.km)
        coord_frame = SkyOffsetFrame(origin=coord)
        if dist:
            off_ra = np.arctan2(off_ra, distance)
            off_dec = np.arctan2(off_dec, distance)
        if e_dist:
            e_off_ra = np.arctan2(e_off_ra, distance)
            e_off_dec = np.arctan2(e_off_dec, distance)
        new_pos = SkyCoord(lon=off_ra, lat=off_dec, frame=coord_frame)
        new_pos = new_pos.icrs

        error_star = self.star.error_at(self.tca)
        error_ra = np.sqrt(error_star[0]**2 + e_off_ra**2)
        error_dec = np.sqrt(error_star[1]**2 + e_off_dec**2)

        out = 'Ephemeris offset (km): X = {:.1f} +/- {:.1f}; Y = {:.1f} +/- {:.1f}\n'.format(
            distance*np.sin(off_ra.to(u.mas)).value, distance*np.sin(e_off_ra.to(u.mas)).value,
            distance*np.sin(off_dec.to(u.mas)).value, distance*np.sin(e_off_dec.to(u.mas)).value)
        out += 'Ephemeris offset (mas): da_cos_dec = {:.3f} +/- {:.3f}; d_dec = {:.3f} +/- {:.3f}\n'.format(
            off_ra.to(u.mas).value, e_off_ra.to(u.mas).value, off_dec.to(u.mas).value, e_off_dec.to(u.mas).value)
        out += '\nAstrometric object position at time {} for reference {}\n'.format(time.iso, observer.__repr__())
        out += 'RA = {} +/- {:.3f} mas; DEC = {} +/- {:.3f} mas'.format(
            new_pos.ra.to_string(u.hourangle, precision=7, sep=' '), error_ra.to(u.mas).value,
            new_pos.dec.to_string(u.deg, precision=6, sep=' '), error_dec.to(u.mas).value)

        if verbose:
            print(out)
        else:
            return out

    # remove this block for v1.0
    @deprecated_function(message="Please use chords.plot_chords to have a better control of the plots")
    def plot_chords(self, all_chords=True, positive_color='blue', negative_color='green', error_color='red',
                    ax=None, lw=2):
        """Plots the chords of the occultation.

        Parameters
        ----------
        all_chords : `bool`, default=True
            If True, it plots all the chords. If False, it sees what was
            deactivated in self.positions and ignores them.

        positive_color : `str`, default='blue'
            Color for the positive chords.

        negative_color : `str`, default='green'
            Color for the negative chords.

        error_color : `str`, default='red'
            Color for the error bars of the chords.

        ax : `maptlotlib.pyplot.Axes`, default=None
            Axis where to plot chords (default: Uses matplotlib pool).

        lw : `int`, `float`, default=2
            Linewidth of the chords.

        """
        self.chords.plot_chords(segment='positive', only_able=not all_chords, color=positive_color, lw=lw, ax=ax)
        self.chords.plot_chords(segment='error', only_able=not all_chords, color=error_color, lw=lw, ax=ax)
        self.chords.plot_chords(segment='negative', color=negative_color, lw=lw, ax=ax, linestyle='--')

    # end of block removal

    def get_map_sites(self):
        """Returns Dictionary with sites in the format required by plot_occ_map function.

        Returns
        -------
        sites : `dict`
            Dictionary with the sites in the format required by `plot_occ_map`
            function.
        """
        sites = {}
        color = {'positive': 'blue', 'negative': 'red'}
        for name, chord in self.chords.items():
            obs = chord.observer
            sites[name] = [obs.lon.deg, obs.lat.deg, 10, 10, color[chord.status()]]
        return sites

    def plot_occ_map(self, **kwargs):
        if 'radius' not in kwargs and hasattr(self, 'fitted_params'):
            r_equa = self.fitted_params['equatorial_radius'][0]
            obla = self.fitted_params['oblateness'][0]
            pos_ang = self.fitted_params['position_angle'][0]
            theta = np.linspace(-np.pi, np.pi, 1800)
            map_pa = self.predict['P/A'][0]
            circle_x = r_equa*np.cos(theta)
            circle_y = r_equa*(1.0-obla)*np.sin(theta)
            ellipse_y = -circle_x*np.sin((pos_ang-map_pa)*u.deg) + circle_y*np.cos((pos_ang-map_pa)*u.deg)
            kwargs['radius'] = ellipse_y.max()
            print('Projected shadow radius = {:.1f} km'.format(kwargs['radius']))
        kwargs['sites'] = kwargs.get('sites', self.get_map_sites())
        if 'offset' not in kwargs and hasattr(self, 'fitted_params'):
            off_ra = self.fitted_params['center_f'][0]*u.km
            off_dec = self.fitted_params['center_g'][0]*u.km
            off_ra = np.arctan2(off_ra, self.dist)
            off_dec = np.arctan2(off_dec, self.dist)
            kwargs['offset'] = [off_ra, off_dec]
        self.predict.plot_occ_map(**kwargs)

    plot_occ_map.__doc__ = PredictionTable.plot_occ_map.__doc__

    def to_log(self, namefile=None):
        """Saves the occultation log to a file.

        Parameters
        ----------
        namefile : `str`
            Filename to save the log.
        """
        if namefile is None:
            namefile = 'occ_{}_{}.log'.format(self.body.shortname.replace(' ', '_'), self.tca.isot[:16])
        f = open(namefile, 'w')
        f.write(self.__str__())
        f.close()

    def check_time_shift(self, time_interval=30, time_resolution=0.001, verbose=False, plot=False, use_error=True, delta_plot=100,
                         ignore_chords=None):
        """Check the needed time offset, so all chords have their center aligned.

        Parameters
        ----------
        time_interval : `int`, `float`
            Time interval to check, default is 30 seconds.

        time_resolution : `int`, `float`
            Time resolution of the search, default is 0.001 seconds.

        verbose : `bool`
            If True, it prints text, default is False.

        plot : `bool`, default=False
            If True, it plots figures as a visual aid.

        use_error : `bool`, default=True
            if True, the linear fit considers the time uncertainty.

        delta_plot : `int`, `float`, default=100
            Value to be added to increase the plot limit, in km.

        ignore_chords : `str`, list, default=None
            Names of the chords to be ignored in the linear fit.

        Returns
        -------
        time_shift : `dict`
            Dictionary with needed time offset to align the chords, each key is
            the name of the chord.
        """
        import matplotlib.pyplot as plt
        fm = np.array([])
        gm = np.array([])
        dfm = np.array([])
        dgm = np.array([])
        vfm = np.array([])
        vgm = np.array([])
        chord_name = np.array([])
        ignore_chords = np.array(ignore_chords, ndmin=1)
        use_chords = np.array([], dtype=bool)
        out_dic = {}

        chords = range(len(self.chords))
        for i in chords:
            chord = self.chords[i]
            if chord.status() == 'positive':
                ffm, ggm, vffm, vggm = chord.get_fg(time=chord.lightcurve.time_mean, vel=True)
                df1, dg1, df2, dg2 = chord.path(segment='error')
                dfm = np.append(dfm, np.sqrt(((df1.max() - df1.min())/2)**2 + ((df2.max() - df2.min())/2)**2))
                dgm = np.append(dgm, np.sqrt(((dg1.max() - dg1.min())/2)**2 + ((dg2.max() - dg2.min())/2)**2))
                fm = np.append(fm, ffm)
                gm = np.append(gm, ggm)
                vfm = np.append(vfm, vffm)
                vgm = np.append(vgm, vggm)
                chord_name = np.append(chord_name, chord.name)
                if chord.name in ignore_chords:
                    use_chords = np.append(use_chords, False)
                else:
                    use_chords = np.append(use_chords, True)
        if len(fm[use_chords]) < 2:
            raise ValueError('The number of fitted chords should be higher than two')
        out = self.__linear_fit_error(x=fm[use_chords], y=gm[use_chords], sx=dfm[use_chords], sy=dgm[use_chords],
                                      verbose=verbose, use_error=use_error)
        fm_fit = np.arange(-delta_plot+np.min([fm.min(), gm.min()]), delta_plot+np.max([fm.max(), gm.max()]),
                           time_resolution*np.absolute(self.vel.value))
        gm_fit = self.__func_line_decalage(out.beta, fm_fit)

        previous_dt = np.array([])
        fm_decalage = np.array([])
        gm_decalage = np.array([])
        time_decalage = np.array([])

        for i, j in enumerate(chord_name):
            previous_dt = np.append(previous_dt, self.chords[j].lightcurve.dt)
            dist_min = np.array([])
            dtt = np.arange(-time_interval, +time_interval, time_resolution)
            for dt in dtt:
                fm_new = fm[i] + dt*vfm[i]
                gm_new = gm[i] + dt*vgm[i]
                dist = np.sqrt((fm_new - fm_fit)**2 + (gm_new - gm_fit)**2)
                dist_min = np.append(dist_min, dist.min())
            print('dt = {:+6.3f} seconds; chord = {}'.format(previous_dt[i] + dtt[dist_min.argmin()], chord_name[i]))
            fm_decalage = np.append(fm_decalage, fm[i] + dtt[dist_min.argmin()]*vfm[i])
            gm_decalage = np.append(gm_decalage, gm[i] + dtt[dist_min.argmin()]*vgm[i])
            time_decalage = np.append(time_decalage, dtt[dist_min.argmin()])

        for i in range(len(chord_name)):
            out_dic[chord_name[i]] = time_decalage[i]
        if plot:
            plt.figure(figsize=(5, 5))
            plt.title('Before', fontsize=15)
            self.chords.plot_chords(color='blue')
            self.chords.plot_chords(color='red', segment='error')
            plt.plot(fm[use_chords], gm[use_chords], linestyle='None', marker='o', color='k')
            plt.plot(fm[np.invert(use_chords)], gm[np.invert(use_chords)], linestyle='None', marker='x', color='r')
            plt.plot(fm_fit, gm_fit, 'k-')
            plt.xlim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.ylim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.show()
            for i, j in enumerate(chord_name):
                self.chords[j].lightcurve.dt = previous_dt[i] + time_decalage[i]
            plt.figure(figsize=(5, 5))
            plt.title('After', fontsize=15)
            self.chords.plot_chords(color='blue')
            self.chords.plot_chords(color='red', segment='error')
            plt.plot(fm_decalage[use_chords], gm_decalage[use_chords], linestyle='None', marker='o', color='k')
            plt.plot(fm_decalage[np.invert(use_chords)], gm_decalage[np.invert(use_chords)], linestyle='None', marker='x',
                     color='r')
            plt.plot(fm_fit, gm_fit, 'k-')
            plt.xlim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.ylim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.show()
            for i, j in enumerate(chord_name):
                self.chords[j].lightcurve.dt = previous_dt[i]
        return out_dic

    def to_file(self):
        """Saves the occultation data to a file.

        Three files are saved containing the positions and velocities for the
        observations. They are for the positive, negative and error bars positions.

        The format of the files are: positions in f and g, velocities in f and
        g, the Julian Date of the observation, light curve name of the
        corresponding position.
        """
        pos = []
        neg = []
        err = []
        for name, chord in self.chords.items():
            status = chord.status()
            l_name = name.replace(' ', '_')
            if status == 'positive':
                im = chord.lightcurve.immersion
                ime = chord.lightcurve.immersion_err
                f, g, vf, vg = chord.get_fg(time=im, vel=True)
                f1, g1 = chord.get_fg(time=im-ime*u.s)
                f2, g2 = chord.get_fg(time=im+ime*u.s)
                pos.append([f, g, vf, vg, im.jd, l_name+'_immersion'])
                err.append([f1, g1, vf, vg, (im-ime*u.s).jd, l_name+'_immersion_err-'])
                err.append([f2, g2, vf, vg, (im+ime*u.s).jd, l_name+'_immersion_err+'])

                em = chord.lightcurve.emersion
                eme = chord.lightcurve.emersion_err
                f, g, vf, vg = chord.get_fg(time=em, vel=True)
                f1, g1 = chord.get_fg(time=em-eme*u.s)
                f2, g2 = chord.get_fg(time=em+eme*u.s)
                pos.append([f, g, vf, vg, em.jd, l_name+'_emersion'])
                err.append([f1, g1, vf, vg, (em-eme*u.s).jd, l_name+'_emersion_err-'])
                err.append([f2, g2, vf, vg, (em+eme*u.s).jd, l_name+'_emersion_err+'])

            if status == 'negative':
                ini = chord.lightcurve.initial_time
                f, g, vf, vg = chord.get_fg(time=ini, vel=True)
                neg.append([f, g, vf, vg, ini.jd, l_name+'_start'])

                end = chord.lightcurve.end_time
                f, g, vf, vg = chord.get_fg(time=end, vel=True)
                neg.append([f, g, vf, vg, end.jd, l_name+'_end'])
        if len(pos) > 0:
            f = open('occ_{}_pos.txt'.format(self.body.shortname.replace(' ', '_')), 'w')
            for line in pos:
                f.write('{:10.3f} {:10.3f} {:-6.2f} {:-6.2f} {:16.8f} {}\n'.format(*line))
            f.close()
            f = open('occ_{}_err.txt'.format(self.body.shortname.replace(' ', '_')), 'w')
            for line in err:
                f.write('{:10.3f} {:10.3f} {:-6.2f} {:-6.2f} {:16.8f} {}\n'.format(*line))
            f.close()
        if len(neg) > 0:
            f = open('occ_{}_neg.txt'.format(self.body.shortname.replace(' ', '_')), 'w')
            for line in neg:
                f.write('{:10.3f} {:10.3f} {:-6.2f} {:-6.2f} {:16.8f} {}\n'.format(*line))
            f.close()

    def __str__(self):
        """ String representation of the Occultation class
        """
        out = ('Stellar occultation of star {} {} by {}.\n\n'
               'Geocentric Closest Approach: {:.3f}\n'
               'Instant of CA: {}\n'
               'Position Angle: {:.2f}\n'
               'Geocentric shadow velocity: {:.2f}\n'
               'Sun-Geocenter-Target angle:  {:.2f} deg\n'
               'Moon-Geocenter-Target angle: {:.2f} deg\n\n\n'.format(
            self.star._catalogue, self.star.code, self.body.name, self.ca, self.tca.iso,
            self.pa, self.vel, self.predict['S-G-T'].data[0], self.predict['M-G-T'].data[0])
        )

        count = {'positive': 0, 'negative': 0, 'visual': 0}
        string = {'positive': '', 'negative': '', 'visual': ''}
        if len(self.chords) > 0:
            for chord in self.chords.values():
                status = chord.status()
                string[status] += chord.__str__() + '\n'
                count[status] += 1

        if len(self.chords) == 0:
            out += 'No observations reported'
        else:
            out += '\n'.join(['{} {} observations'.format(count[k], k) for k in string.keys() if count[k] > 0])

        out += '\n\n'

        out += '#'*79 + '\n{:^79s}\n'.format('STAR') + '#'*79 + '\n'
        out += self.star.__str__() + '\n\n'

        coord = self.star.get_position(self.tca, observer=self._reference_center)
        try:
            error_star = self.star.error_at(self.tca)
        except:
            error_star = [0, 0]*u.mas
        out += 'Geocentric star coordinate at occultation Epoch ({}):\n'.format(self.tca.iso)
        out += 'RA={} +/- {:.4f}, DEC={} +/- {:.4f}\n\n'.format(
            coord.ra.to_string(u.hourangle, sep='hms', precision=5), error_star[0],
            coord.dec.to_string(u.deg, sep='dms', precision=4), error_star[1])

        out += self.body.__str__() + '\n'

        for status in string.keys():
            if count[status] == 0:
                continue
            out += ('#'*79 + '\n{:^79s}\n'.format(status.upper() + ' OBSERVATIONS') + '#'*79)
            out += '\n\n'
            out += string[status]

        if hasattr(self, 'fitted_params'):
            out += '#'*79 + '\n{:^79s}\n'.format('RESULTS') + '#'*79 + '\n\n'
            out += 'Fitted Ellipse:\n'
            out += '\n'.join(['{}: {:.3f} +/- {:.3f}'.format(k, *self.fitted_params[k])
                              for k in self.fitted_params.keys()]) + '\n'
            polar_radius = self.fitted_params['equatorial_radius'][0]*(1.0-self.fitted_params['oblateness'][0])
            equivalent_radius = np.sqrt(self.fitted_params['equatorial_radius'][0]*polar_radius)
            out += 'polar_radius: {:.3f} km \n'.format(polar_radius)
            out += 'equivalent_radius: {:.3f} km \n'.format(equivalent_radius)
            if not np.isnan(self.body.H):
                H_sun = -26.74
                geometric_albedo = (10**(0.4*(H_sun - self.body.H.value))) * ((u.au.to('km')/equivalent_radius)**2)
                out += 'geometric albedo (V): {:.3f} ({:.1%}) \n'.format(geometric_albedo, geometric_albedo)
            else:
                out += 'geometric albedo (V): not calculated, absolute magnitude (H) is unknown \n'
            out += '\nMinimum chi-square: {:.3f}\n'.format(self.chi2_params['chi2_min'])
            out += 'Number of fitted points: {}\n'.format(self.chi2_params['npts'])
            out += 'Number of fitted parameters: {}\n'.format(self.chi2_params['nparam'])
            out += 'Minimum chi-square per degree of freedom: {:.3f}\n'.format(
                self.chi2_params['chi2_min']/(self.chi2_params['npts'] - self.chi2_params['nparam']))
            out += 'Radial dispersion: {:.3f} +/- {:.3f} km\n'.format(
                self.chi2_params['radial_dispersion'].mean(), self.chi2_params['radial_dispersion'].std(ddof=1))
            out += 'Radial error:      {:.3f} +/- {:.3f} km\n'.format(
                self.chi2_params['radial_error'].mean(), self.chi2_params['radial_error'].std(ddof=1))

            out += '\n' + self.new_astrometric_position(verbose=False)

        return out

    def __func_line_decalage(self, p, x):
        """ Private function returns a linear function, intended for fitting inside the self.check_decalage().
        """
        a, b = p
        return a*x + b

    def __linear_fit_error(self, x, y, sx, sy, verbose=False, use_error=True):
        """ Private function returns a linear fit, intended for fitting inside the self.check_decalage().
        """
        import scipy.odr as odr

        model = odr.Model(self.__func_line_decalage)
        if use_error:
            data = odr.RealData(x=x, y=y, sx=sx, sy=sy)
        else:
            data = odr.RealData(x=x, y=y)
        fit = odr.ODR(data, model, beta0=[0., 1.])
        out = fit.run()
        if verbose:
            print('Linear fit procedure')
            out.pprint()
            print('\n')
        return out

import numpy,import scipy,import astropy,from astropy


import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

#%%

#f=open('lane1_glitches.dat', mode='rb')
f=open('fastGlitches1.dat', mode='rb')

data = f.read()
f.close()
len(data)

#%%

################################################################
# Print seleced POD header
################################################################

#header = np.frombuffer(data, dtype='uint16', count=12, offset=(192+12)*2)
#for i in range(0, 12):
#   print('%d %04x' %(i, header[i]))

#%%

dataOffset = 8
podNo = 0

while 1:
   header = np.frombuffer(data, dtype='uint16', count=12, offset=dataOffset)

   adcCh    = header[2] & 0xFF
   if ((header[3] & 0xF000) >> 12) == 0:
      adcType = 'fast'
   else:
      adcType = 'slow'
   debugInfo= header[1]
   isFooter = header[3] & 0x1
   trigSize = ((header[5] & 0x3F) << 16) | header[4]
   trigOfs  = (header[7] << 16) | header[6]
   fastOfs  = (header[5] & 0x180) >> 7
   lostFlg  = (header[5] & 0x40) >> 6
   extFlg   = (header[5] & 0x800) >> 11
   intFlg   = (header[5] & 0x1000) >> 12
   emptyFlg = (header[5] & 0x2000) >> 13
   vetoFlg  = (header[5] & 0x4000) >> 14
   badAFlg  = (header[5] & 0x8000) >> 15
   trigTime = (header[11] << 48) | (header[10] << 32) | (header[9] << 16) | header[8]

   adcData = np.frombuffer(data, dtype='uint16', count=trigSize, offset=dataOffset+24)
   #print('############ POD number: %d ############' %(podNo))
   #print('ADC max %d' %(max(adcData)))
   #print('ADC min %d' %(min(adcData)))

   if max(adcData) > 39000 or min(adcData) < 27000:
      print('############ POD number: %d ############' %(podNo))
      print('ADC channel %d' %(adcCh))
      print('ADC type %s' %(adcType))
      print('Footer present %d' %(isFooter))
      print('Trig size %d' %(trigSize))
      print('Trig offset %d' %(trigOfs))
      print('Fast ADC sample offset %d' %(fastOfs))
      print('Lost flag %d' %(lostFlg))
      print('Ext flag %d' %(extFlg))
      print('Int flag %d' %(intFlg))
      print('Empty flag %d' %(emptyFlg))
      print('Veto flag %d' %(vetoFlg))
      print('Bad ADC flag %d' %(badAFlg))
      print('Header time %d' %(trigTime))
      if adcType == 'fast':
         print('Debug info %s' %(format(debugInfo, 'b')))
      print('ADC max %d' %(max(adcData)))
      print('ADC min %d' %(min(adcData)))
      #save data in csv file
      fileName = 'POD' + str(podNo) +'_ch' + str(adcCh) + adcType + '_T' + str(trigTime)

      if podNo == 741:
          break

      #plt.plot(adcData, '.')
      #plt.savefig(fileName + '.png')
      #plt.clf()
      #plt.close()

      #f = open(fileName + '.csv', 'w')
      #for i in range(len(adcData)):
      #   f.write('%d\n' %(adcData[i]))
      #f.write('\n')
      #f.close()

   #check for bad ADC flag
   if badAFlg == 1:
      print('########################## Bad ADC flag in POD %d' %(podNo))

   # correct odd trigSize
   if trigSize%2 != 0:
      trigSize = trigSize + 1

   #move to next offset
   dataOffset = dataOffset + (trigSize+12)*2 + 8

   podNo = podNo + 1

   if dataOffset > len(data):
      print("No more data for POD number %d" %(podNo+1))
      break

#%%

fig = plt.figure(1,figsize=(12,8),dpi=150)
plt.plot(bins)
print(len(adcData))

#%%

bins = 2**16 * [0]
for i in range(len(adcData)):
    bins[adcData[i]] = bins[adcData[i]] + 1

#%%

for i in range(2**16):
    if bins[i] != 0:
        print('bin %d val %d'%(i, bins[i]))


import os
import random
from itertools import chain
from pathlib import Path

import numpy as np
import pytest
from jina.executors.indexers import BaseIndexer

from jina import Document
from jina.flow import Flow

TOP_K = 10

@pytest.fixture
def config(tmpdir):
    random.seed(0)
    np.random.seed(0)
    os.environ['JINA_CRUD_CHUNKS'] = str(tmpdir)
    os.environ['JINA_TOPK'] = str(TOP_K)
    yield
    del os.environ['JINA_CRUD_CHUNKS']
    del os.environ['JINA_TOPK']

d_embedding = np.random.random([9])
c_embedding = np.random.random([9])

def document_generator(start, num_docs, num_chunks):
    chunk_id = num_docs
    for idx in range(start, num_docs):
        with Document() as doc:
            doc.id = idx
            doc.tags['id'] = idx
            doc.text = f'I have {idx} cats'
            doc.embedding = np.random.random([9])
            for chunk_idx in range(num_chunks):
                with Document() as chunk:
                    chunk.id = chunk_id
                    chunk.tags['id'] = chunk_idx
                    chunk.text = f'I have {chunk_idx} chunky cats'
                    chunk.embedding = np.random.random([9])
                chunk_id += 1
                doc.chunks.append(chunk)
        yield doc

def validate_index_size(num_indexed_docs):
    path = Path(os.environ['JINA_CRUD_CHUNKS'])
    index_files = list(path.glob('*.bin'))
    assert len(index_files) > 0
    for index_file in index_files:
        index = BaseIndexer.load(str(index_file))
        assert index.size == num_indexed_docs

@pytest.mark.parametrize('flow_file', ['flow_vector.yml'])
def test_delete_vector(config, mocker, flow_file):
    num_searches = 10
    num_docs = 10
    num_chunks = 5

    def validate_result_factory(num_matches):
        def validate_results(resp):
            mock()
            assert len(resp.docs) == num_searches
            for doc in resp.docs:
                assert len(doc.matches) == num_matches

        return validate_results

    with Flow.load_config(flow_file) as index_flow:
        index_flow.index(
            input_fn=document_generator(start=0, num_docs=num_docs, num_chunks=num_chunks))
    validate_index_size(num_chunks*num_docs) #5 chunks for each of the 10 docs

    mock = mocker.Mock()
    with Flow.load_config(flow_file) as search_flow:
        search_flow.search(
            input_fn=document_generator(start=0, num_docs=num_docs, num_chunks=num_chunks),
            on_done=validate_result_factory(TOP_K))
    mock.assert_called_once()

    with Flow.load_config(flow_file) as index_flow:
        index_flow.delete(
            input_fn=document_generator(start=0, num_docs=num_docs, num_chunks=num_chunks))
    validate_index_size(0)

    mock = mocker.Mock()
    with Flow.load_config(flow_file) as search_flow:
        search_flow.search(
            input_fn=document_generator(start=0, num_docs=num_docs, num_chunks=num_chunks),
            on_done=validate_result_factory(0))
    mock.assert_called_once()

@pytest.mark.parametrize('flow_file', ['flow_vector.yml'])
def test_update_vector(config, mocker, flow_file):
    num_searches = 10
    num_docs = 10
    num_chunks = 5

    docs_before = list(document_generator(start=0, num_docs=num_docs, num_chunks=num_chunks))
    docs_updated = list(document_generator(start=10, num_docs=20, num_chunks=num_chunks))
    ids_before = list()
    ids_updated = list()

    def validate_result_factory(has_changed, num_matches):
        def validate_results(resp):
            mock()
            assert len(resp.docs) == num_searches
            for d in docs_before:
                ids_before.append(d.id)
            for d in docs_updated:
                ids_updated.append(d.id)
            for doc in resp.docs:
                assert len(doc.matches) == num_matches
                if has_changed:
                    assert doc.id in ids_updated
                    assert doc.id not in ids_before
                else:
                    assert doc.id in ids_before
                    assert doc.id not in ids_updated
        return validate_results

    with Flow.load_config(flow_file) as index_flow:
        index_flow.index(input_fn=docs_before)
    validate_index_size(num_chunks*num_docs) #num_docs per all its chunks, 50 in this case

    mock = mocker.Mock()
    with Flow.load_config(flow_file) as search_flow:
        search_flow.search(
            input_fn=document_generator(start=0, num_docs=num_docs, num_chunks=num_chunks),
            on_done=validate_result_factory(has_changed=False, num_matches=TOP_K))
    mock.assert_called_once()

    with Flow.load_config(flow_file) as index_flow:
        index_flow.update(input_fn=docs_updated)
    validate_index_size(num_chunks*num_docs) #num_docs per all its chunks, 50 in this case

    mock = mocker.Mock()
    with Flow.load_config(flow_file) as search_flow:
        search_flow.search(
            input_fn=document_generator(start=10, num_docs=20, num_chunks=num_chunks),
            on_done=validate_result_factory(has_changed=True, num_matches=num_docs))
    mock.assert_called_once()