#!/usr/bin/env python
"""
Building blocks for Blow

Serra, J., Pascual, S. & Segura, C. Blow: a single-scale hyperconditioned flow
for non-parallel raw-audio voice conversion. in Proc. NIPS (2019).

Reference: https://github.com/joansj/blow
"""
from __future__ import absolute_import

import sys
import numpy as np

import torch
import torch.nn as torch_nn
import torch.nn.functional as torch_nn_func
import torch.nn.init as torch_init

import sandbox.block_glow as nii_glow
import core_scripts.data_io.wav_tools as nii_wav_tk
import core_scripts.data_io.conf as nii_io_conf
import core_scripts.other_tools.debug as nii_debug

__author__ = "Xin Wang"
__email__ = "wangxin@nii.ac.jp"
__copyright__ = "Copyright 2021, Xin Wang"

#######################################
# Numpy utilities for data augmentation
#######################################
def flip(x):
    """y=flip(x) flips the sign of x
    input: x, np.array
    output: y, np.array
    """
    return np.sign(np.random.rand(1)-0.5) * x

def ampscale(x):
    """y=ampscale(x) randomly scale the amplitude of x
    input: x, np.array
    output: y, np.array
    """
    return (2*np.random.rand(1)-1) * x / (np.max(np.abs(x)) + 1e-07)

def framejitter(x, framelen):
    """y=framejitter(x, framelen)
    input: x, np.array, original waveform (length, 1)
           framelen, int, framelen
    output: y, np.array, segment of the waveform
    """
    framelen = x.shape[0] if framelen > x.shape[0] else framelen
    random_start = int(np.ceil(np.random.rand(1) * (x.shape[0] - framelen)))
    return x[random_start:random_start+framelen]

def emphasis_rand(x, coef_val):
    """y=deemphasis(x, coef_val)
    input: x, np.array, original waveform (length, 1) or (length)
           framelen, int, framelen
    output: y, np.array, segment of the waveform
    """
    coef = (2 * np.random.rand(1) - 1) * coef_val
    x_new = np.zeros_like(x) + x
    x_new[1:] = x_new[1:] - coef * x[:-1]
    return x_new

def wav_aug(x, framelen, coef_val, sr):
    """y = wav_aug(x, framelen, coef_val, sr)
    input
    -----
      x: np.array, original waveform (length, 1)
      framelen: int, frame length
      coef_val: float, reference coefficient for emphasis-rand
      sr: int, sampling rate (e.g., 16000)

    output
    ------
      y: np.array, pre-processed waveform (length, 1)
    """
    trimmed_x = nii_wav_tk.silence_handler_wrapper(x, sr, flag_output=1)
    x_frame = framejitter(trimmed_x, framelen)
    return ampscale(emphasis_rand(x_frame, coef_val))

class OverlapAdder(torch_nn.Module):
    """OverlapAdder
    """
    def __init__(self, fl, fs, flag_win_analysis=True):
        """OverlapAdder(flag_windowing_before=True)
        Args
        ----
          fl: int, frame length
          fs: int, frame shift
          flag_win_analysis: bool (default True)
              True: apply windowing during analysis
              False: apply windowing during synthesis
        """
        super(OverlapAdder, self).__init__()
        self.fl = fl
        self.fs = fs
        self.flag_win_ana = flag_win_analysis

        # assume even
        self.m_win = torch_nn.Parameter(torch.hann_window(self.fl))
        return

    def get_frame_num(self, wav_length):
        """frame_num = get_frame_num(wav_length)
        wav_length: int, waveform length
        frame_num: int, number of frames
        """
        return (wav_length - self.fl) // self.fs + 1

    def get_wavlength(self, frame_num):
        """wav_length = get_wavlength(self, frame_num)
        wav_length: int, waveform length
        frame_num: int, number of frames
        """
        return (frame_num - 1) * self.fs + self.fl

    def forward(self, x):
        """OverlapAdder(x)

        input
        -----
          x: tensor, (batch, length, 1)

        output
        ------
          y: tensor, (batch, frame_num, frame_length)
        """
        frame_num = self.get_frame_num(x.shape[1])

        # (batch, num_patches, 1, patch_size)
        # num_patches = (length - length) // shift + 1
        # and copy the data
        # note that unfold put each patch as the last dimension
        # x_tmp (batch, frame_num, 1, frame_length)
        x_tmp = x.unfold(1, self.fl, self.fs)

        # apply window
        if self.flag_win_ana:
            x_tmp = x_tmp * self.m_win

        # (batch, frame_num, frame_length)
        return x_tmp.view(x.shape[0], x_tmp.shape[1], -1)

    def reverse(self, x_framed, flag_scale=False):
        """OverlapAdder(x)

        input
        -----
          x: tensor, (batch, frame_num, frame_length)
          flag_scale: bool, whether scale the ampltidue to (-1, 1)
                      default False
        output
        ------
          y: tensor, (batch, length, 1)
        """
        batch, frame_num, frame_len = x_framed.shape
        x_len = self.get_wavlength(frame_num)
        x_buf = torch.zeros(
            [batch, x_len], device=x_framed.device, dtype=x_framed.dtype)
        x_win = torch.zeros_like(x_buf)

        for idx in range(frame_num):
            sdx = idx * self.fs
            edx = sdx + self.fl
            x_win[:, sdx:edx] += self.m_win
            if not self.flag_win_ana:
                x_buf[:, sdx:edx] += x_framed[:, idx] * self.m_win
            else:
                x_buf[:, sdx:edx] += x_framed[:, idx]
        # assume the overlapped window has a constant amplitude
        x_buf = x_buf / x_win.mean()

        # normalize the amplitude between (-1, 1)
        if flag_scale:
            # if input is between (-1, 1), there is no need to
            # do this normalization
            x_buf = x_buf / (x_buf.abs().max())
        return x_buf.unsqueeze(-1)

#######################################
# Torch model definition
#######################################

class AffineCouplingBlow_core(torch_nn.Module):
    """AffineCouplingBlow_core

    AffineCoupling core layer the produces the scale and bias parameters.

    Example:
        feat_dim = 10
        cond_dim = 20

        m_layer = AffineCouplingBlow_core(feat_dim, cond_dim, 64, 2)

        data = torch.randn([2, 100, feat_dim])
        cond = torch.randn([2, 1, cond_dim])
        scale, bias, log_scale = m_layer(data, cond)
    """
    def __init__(self, feat_dim, cond_dim, num_ch, kernel_size=3):
        """AffineCouplingBlow_core(feat_dim, cond_dim, num_ch, kernel_size=3)

        Args
        ----
          feat_dim: int, dimension of input feature
          cond_dim: int, dimension of conditional features
          num_ch: int, number of channels for conv layers
          kernel_size: int, kernel size of conv layer, default 3

        input_feature -------> func.conv1d -----> conv1ds -> scale, bias
                                    ^
                                    |
        cond_dim ---> Adapter -> conv weight/bias
        """
        super(AffineCouplingBlow_core, self).__init__()

        self.feat_dim = feat_dim
        self.cond_dim = cond_dim

        # make sure that kernel is odd
        if kernel_size % 2 == 0:
            self.kernel_s = kernel_size + 1
            print("\tAffineCouplingBlow_core", end=" ")
            print("kernel size {:d} -> {:d}".format(kernel_size, self.kernel_s))
        else:
            self.kernel_s = kernel_size

        if num_ch % feat_dim != 0:
            # make sure that number of channel is good
            self.num_ch = num_ch // feat_dim * feat_dim
            print("\tAffineCouplingBlow_core", end=" ")
            print("conv channel {:d} -> {:d}".format(num_ch, self.num_ch))
        else:
            self.num_ch = num_ch

        # Adapter
        # (batch, 1, cond_dim) -> (batch, 1, kernel_size * num_ch) for weight
        #                      -> (batch, 1, num_ch) for bias
        self.m_adapter = torch_nn.Linear(cond_dim,
                                         (self.kernel_s+1) * self.num_ch)

        # conv1d with condition-independent parameters
        self.m_conv1ds = torch_nn.Sequential(
            torch_nn.ReLU(),
            torch_nn.Conv1d(self.num_ch, self.num_ch, 1),
            torch_nn.ReLU(),
            torch_nn.Conv1d(self.num_ch, feat_dim * 2, self.kernel_s,
                           padding=(self.kernel_s-1)//2)
        )

        # zero initialization for the last conv layers
        # similar to Glow and WaveGlow
        self.m_conv1ds[-1].weight.data.zero_()
        self.m_conv1ds[-1].bias.data.zero_()
        return

    def forward(self, x, cond):
        """scale, bias = AffineCouplingBlow_core(x, cond)

        input
        -----
          x: tensor, input tensor (batch, length, feat_dim)
          cond: tensor, condition feature (batch, 1, cond_dim)

        output
        ------
          scale: tensor, scaling parameters (batch, length, feat_dim)
          bias: tensor, bias paramerters (batch, length, feat_dim)
        """
        # cond_dim -> Adapter -> conv weight/bias
        # cond[:, 0, :] -> (batch, cond_dim)
        # adapter(cond[:, 0, :]) -> (batch, kernel_size * num_ch + num_ch)
        # view(...) -> (batch * num_ch, kernel_size + 1)
        weight_bias = self.m_adapter(cond[:, 0, :]).view(-1, self.kernel_s+1)
        # (batch * num_ch, 1, kernel_size)
        weight = weight_bias[:, 0:self.kernel_s].unsqueeze(1)
        # (batch * num_ch)
        bias = weight_bias[:, self.kernel_s]

        # convolution given weight_bias
        padsize = (self.kernel_s - 1) // 2
        groupsize = x.shape[0] * self.feat_dim
        length = x.shape[1]

        #  x.permute(0, 2, 1)...view -> (1, batch*feat_dim, length)
        #  conv1d -> (1, batch * num_ch, length)
        #  view -> (batch, num_ch, length)
        x_tmp = torch_nn_func.conv1d(
            x.permute(0, 2, 1).contiguous().view(1, -1, length),
            weight,
            bias = bias,
            padding = padsize,
            groups = groupsize
        ).view(x.shape[0], -1, length)

        # condition invariant conv -> (batch, feat_dim * 2, length)
        x_tmp = self.m_conv1ds(x_tmp)

        # scale and bias (batch, feat_dim, length)
        raw_scale, bias = torch.chunk(x_tmp, 2, dim=1)

        #  -> (batch, length, feat_dim)
        bias = bias.permute(0, 2, 1)

        #  re-parameterize
        #   Here we need to add a small number, otherwise, log(scale)
        #   somtime times become -inf during training
        scale = torch.sigmoid(raw_scale + 2).permute(0, 2, 1) * 0.5 + 0.5

        log_scale = torch.log(scale)

        #print("Debug: {:.3f} {:.3f} {:.3f} {:3f}".format(
        #    log_scale.max().item(), log_scale.min().item(),
        #    scale.max().item(), scale.min().item()),
        #    file=sys.stderr)
        return scale, bias, log_scale

class AffineCouplingBlow(torch_nn.Module):
    """AffineCouplingBlow

    AffineCoupling block in Blow

    Example:
        feat_dim = 10
        cond_dim = 20

        m_layer = AffineCouplingBlow(feat_dim, cond_dim,60,3, flag_detjac=True)

        data = torch.randn([2, 100, feat_dim])
        cond = torch.randn([2, 1, cond_dim])
        out, detjac = m_layer(data, cond)

        data_rever = m_layer.reverse(out, cond)

        torch.std(data - data_rever)
    """
    def __init__(self, in_dim, cond_dim,
                 conv_dim_channel, conv_kernel_size,
                 flag_detjac=False):
        """AffineCouplingBlow(in_dim, cond_dim,
            wn_num_conv1d, wn_dim_channel, wn_kernel_size,
            flag_affine=True, flag_detjac=False)

        Args:
        -----
          in_dim: int, dim of input audio data (batch, length, in_dim)
          cond_dim, int, dim of condition feature (batch, length, cond_dim)
          conv_dim_channel: int, dime of the convolution channels
          conv_kernel_size: int, kernel size of the convolution layers
          flag_detjac: bool, whether return the determinant of Jacobian,
                       default False

        y -> split() -> y1, y2 -> concate([y1, (y2+bias) * scale])
        When flag_affine == True, y1 -> H() -> scale, bias
        When flag_affine == False, y1 -> H() -> bias, scale=1
        Here, H() is AffineCouplingBlow_core layer
        """
        super(AffineCouplingBlow, self).__init__()

        self.flag_detjac = flag_detjac

        if in_dim % 2 > 0:
            print("AffineCouplingBlow(feat_dim), feat_dim is an odd number?!")
            sys.exit(1)

        # Convolution block to get scale and bias
        self.m_core = AffineCouplingBlow_core(
            in_dim // 2, cond_dim, conv_dim_channel, conv_kernel_size)

        return

    def _detjac(self, log_scale, factor=1):
        # (batch, dim1, dim2, ..., feat_dim) -> (batch)
        # sum over dim1, ... feat_dim
        return nii_glow.sum_over_keep_batch(log_scale / factor)

    def _nn_trans(self, y1, cond):
        """_nn_trans(self, y1, cond)

        input
        -----
          y1: tensor, input feature, (batch, lengh, input_dim//2)
          cond: tensor, condition feature, (batch, length, cond_dim)

        output
        ------
          scale: tensor, (batch, lengh, input_dim // 2)
          bias: tensor, (batch, lengh, input_dim // 2)
          log_scale: tensor, (batch, lengh, input_dim // 2)

        Affine transformaiton can be done by scale * feature + bias
        log_scale is used for det Jacobian computation
        """
        scale, bias, log_scale = self.m_core(y1, cond)
        return scale, bias, log_scale

    def forward(self, y, cond, factor=1):
        """AffineCouplingBlow.forward(y, cond)

        input
        -----
          y: tensor, input feature, (batch, lengh, input_dim)
          cond: tensor, condition feature , (batch, 1, cond_dim)

        output
        ------
          x: tensor, input feature, (batch, lengh, input_dim)
          detjac: tensor, det of jacobian, (batch,)

        y1, y2 = split(y)
        scale, bias = Conv(y1)
        x2 = y2 * scale + bias or (y2 + bias) * scale
        return [y1, x2]
        """
        # split
        y1, y2 = y.chunk(2, -1)
        scale, bias, log_scale = self._nn_trans(y1, cond)

        # transform
        x1 = y1
        x2 = (y2 + bias) * scale

        # concatenate
        x = torch.cat([x1, x2], dim=-1)
        if self.flag_detjac:
            return x, self._detjac(log_scale, factor)
        else:
            return x

    def reverse(self, x, cond):
        """AffineCouplingBlow.reverse(y, cond)

        input
        -----
          x: tensor, input feature, (batch, lengh, input_dim)
          cond: tensor, condition feature , (batch, 1, cond_dim)

        output
        ------
          y: tensor, input feature, (batch, lengh, input_dim)

        x1, x2 = split(x)
        scale, bias = conv(x1)
        y2 = x2 / scale - bias
        return [x1, y2]
        """
        # split
        x1, x2 = x.chunk(2, -1)
        # reverse transform
        y1 = x1
        scale, bias, log_scale = self._nn_trans(y1, cond)
        y2 = x2 / scale - bias
        return torch.cat([y1, y2], dim=-1)

class SqueezeForBlow(torch_nn.Module):
    """SqueezeForBlow

    Squeeze input feature for Blow.

    Example
        data = torch.randn([2, 10, 3])
        m_sq = SqueezeForBlow()
        data_out = m_sq(data)
        data_rev = m_sq.reverse(data_out)
        torch.std(data_rev - data)
    """
    def __init__(self, mode=1):
        """SqueezeForBlow(mode=1)

        Args
        ----
          mode: int, mode of squeeze, default 1

        Mode 1: squeeze by a factor of 2 as in original paper
        """
        super(SqueezeForBlow, self).__init__()

        self.m_mode = mode

        if self.m_mode == 1:
            self.squeeze_factor = 2
        else:
            print("SqueezeForBlow mode {:d} not implemented".format(mode))
            sys.exit(1)
        return

    def get_expected_squeeze_length(self, orig_length):
        # return expected length after squeezing
        if self.m_mode == 1:
            return orig_length // self.squeeze_factor
        else:
            print("unknown mode for SqueezeForBlow")
            sys.exit(1)

    def get_recovered_length(self, squeezed_length):
        # return original length before squeezing
        if self.m_mode == 1:
            return squeezed_length * self.squeeze_factor
        else:
            print("unknown mode for SqueezeForBlow")
            sys.exit(1)

    def get_squeeze_factor(self):
        # return the configuration for squeezing
        if self.m_mode == 1:
            return self.squeeze_factor
        else:
            print("unknown mode for SqueezeForBlow")
            sys.exit(1)

    def forward(self, x):
        """SqueezeForBlow(x)

        input
        -----
          x: tensor, (batch, length, feat_dim)

        output
        ------
          y: tensor, (batch, length//squeeze_factor, feat_dim*squeeze_factor)

        """
        if self.m_mode == 1:
            # squeeze, the 8 points should be the last dimension
            squeeze_len = self.get_expected_squeeze_length(x.shape[1])
            # trim length first
            trim_len = squeeze_len * self.squeeze_factor
            x_tmp = x[:, 0:trim_len, :]

            # (batch, time//squeeze_size, squeeze_size, dim)
            x_tmp = x_tmp.view(x_tmp.shape[0], squeeze_len,
                               self.squeeze_factor, -1)

            # (batch, time//squeeze_size, dim, squeeze_size)
            x_tmp = x_tmp.permute(0, 1, 3, 2).contiguous()

            # (batch, time//squeeze_size, dim * squeeze_size)
            return x_tmp.view(x_tmp.shape[0], squeeze_len, -1)
        else:
            print("SqueezeForWaveGlow not implemented")
            sys.exit(1)
        return x_squeezed

    def reverse(self, x_squeezed):
        if self.m_mode == 1:
            # (batch, time//squeeze_size, dim * squeeze_size)
            batch, squeeze_len, squeeze_dim = x_squeezed.shape

            # (batch, time//squeeze_size, dim, squeeze_size)
            x_tmp = x_squeezed.view(
                batch, squeeze_len, squeeze_dim // self.squeeze_factor,
                self.squeeze_factor)

            # (batch, time//squeeze_size, squeeze_size, dim)
            x_tmp = x_tmp.permute(0, 1, 3, 2).contiguous()

            # (batch, time, dim)
            x = x_tmp.view(batch, squeeze_len * self.squeeze_factor, -1)
        else:
            print("SqueezeForWaveGlow not implemented")
            sys.exit(1)
        return x

class FlowStepBlow(torch_nn.Module):
    """FlowStepBlow
    One flow step for Blow
    y -> intertical_1x1() -> ActNorm -> AffineCoupling -> x

    Example
        feat_dim = 10
        cond_dim = 20

        m_layer = FlowStepBlow(feat_dim, cond_dim, 60, 3)

        data = torch.randn([2, 100, feat_dim])
        cond = torch.randn([2, 1, cond_dim])
        out, detjac = m_layer(data, cond)

        data_rever = m_layer.reverse(out, cond)

        torch.std(data - data_rever)
    """
    def __init__(self, in_dim, cond_dim, conv_dim_channel, conv_kernel_size):
        """FlowStepBlow(in_dim, cond_dim,
                            conv_dim_channel, conv_kernel_size)

        Args
        ----
          in_dim: int, input feature dim, (batch, length, in_dim)
          cond_dim:, int, conditional feature dim, (batch, length, cond_dim)
          cond_dim_channel: int, dim of the convolution layers
          conv_kernel_size: int, kernel size of the convolution layers

        For cond_dim_channel and conv_kernel_size, see AffineCouplingBlow
        """
        super(FlowStepBlow, self).__init__()

        # Invertible transformation layer
        self.m_invtrans = nii_glow.InvertibleTrans(in_dim, flag_detjac=True)

        # Act norm layer
        self.m_actnorm = nii_glow.ActNorm(in_dim, flag_detjac=True)

        # coupling layer
        self.m_coupling = AffineCouplingBlow(
            in_dim, cond_dim, conv_dim_channel, conv_kernel_size,
            flag_detjac=True)

        return

    def forward(self, y, cond, factor=1):
        """FlowStepBlow.forward(y, cond, factor=1)

        input
        -----
          y: tensor, input feature, (batch, lengh, in_dim)
          cond: tensor, condition feature , (batch, 1, cond_dim)
          factor: int, this is used to divde likelihood, default 1
                  if we directly sum all detjac, they will become very large
                  however, we cannot average them directly on y because y
                  may have a different shape from the actual data y
        output
        ------
          x: tensor, input feature, (batch, lengh, input_dim)
          detjac: tensor, det of jacobian, (batch,)
        """
        # 1x1 transform
        x_tmp, log_det_1 = self.m_invtrans(y, factor)

        # Actnorm
        x_tmp, log_det_2 = self.m_actnorm(x_tmp, factor)

        # coupling
        x_tmp, log_det_3 = self.m_coupling(x_tmp, cond, factor)
        return x_tmp, log_det_1 + log_det_2 + log_det_3

    def reverse(self, x, cond):
        """FlowStepBlow.reverse(y, cond)

        input
        -----
          x: tensor, input feature, (batch, lengh, input_dim)
          cond: tensor, condition feature , (batch, 1, cond_dim)

        output
        ------
          y: tensor, input feature, (batch, lengh, input_dim)
        """
        y_tmp1 = self.m_coupling.reverse(x, cond)
        y_tmp2 = self.m_actnorm.reverse(y_tmp1)
        y_tmp3 = self.m_invtrans.reverse(y_tmp2)
        #print("Debug: {:.3f} {:.3f} {:.3f} {:.3f} {:.3f} {:.3f}".format(
        #    y_tmp1.max().item(), y_tmp1.min().item(),
        #    y_tmp2.max().item(), y_tmp2.min().item(),
        #    y_tmp3.max().item(), y_tmp3.min().item()))
        return y_tmp3

class BlowBlock(torch_nn.Module):
    """BlowBlock
    A BlowBlok includes multiple steps of flow for Blow.

    Each block conducts:
    x -> squeeze -> flow step1 -> ... -> flow step N

    Compared with WaveGlowBlock, this is easier because there is no
    multi-scale structure, no need to split the latent z.

    Example:

    """
    def __init__(self, in_dim, cond_dim, n_flow_steps,
                 conv_dim_channel, conv_kernel_size):
        """BlowBlock(in_dim, cond_dim, n_flow_steps,
                 conv_dim_channel, conv_kernel_size)
        Args
        ----
          in_dim: int, input feature dim, (batch, length, in_dim)
          cond_dim:, int, conditional feature dim, (batch, length, cond_dim)
          n_flow_steps: int, number of flow steps in one block
          conv_dim_channel: int, dim of the conv residual and skip channels
          conv_kernel_size: int, kernel size of the convolution layers

        For conv_dim_channel and conv_kernel_size, see AffineCouplingBlow
        """
        super(BlowBlock, self).__init__()

        # squeeze
        self.m_squeeze = SqueezeForBlow()

        squeezed_feat_dim = in_dim * self.m_squeeze.get_squeeze_factor()

        # flow steps
        tmp_flows = []
        for i in range(n_flow_steps):
            tmp_flows.append(
                FlowStepBlow(
                    squeezed_feat_dim, cond_dim,
                    conv_dim_channel, conv_kernel_size))
        self.m_flows = torch_nn.ModuleList(tmp_flows)

        self.m_out_dim = squeezed_feat_dim
        return

    def get_out_feat_dim(self):
        return self.m_out_dim

    def get_expected_squeeze_length(self, orig_length):
        return self.m_squeeze.get_expected_squeeze_length(orig_length)

    def forward(self, y, cond, factor=1):
        """z, log_detjac = BlowBlock(y)

        y -> squeeze -> H() -> z, log_det_jacobian
        H() consists of multiple flow steps (1x1conv + Actnorm + AffineCoupling)

        input
        -----
          y: tensor, (batch, length, dim)
          cond, tensor, (batch, 1, cond_dim)
          factor, None or int, this is used to divde likelihood, default 1

        output
        ------
         log_detjac: tensor or scalar
         z: tensor, (batch, length, dim), for N(z; 0, I) or next flow block
        """
        # squeeze
        x_tmp = self.m_squeeze(y)

        # flows
        log_detjac = 0
        for idx, l_flow in enumerate(self.m_flows):
            x_tmp, log_detjac_tmp = l_flow(x_tmp, cond, factor)
            log_detjac = log_detjac + log_detjac_tmp

        return x_tmp, log_detjac

    def reverse(self, z, cond):
        """y = BlowBlock.reverse(z, cond)

        z -> H^{-1}() -> unsqueeze -> y

        input
        -----
          z: tensor, (batch, length, in_dim)
          cond, tensor, (batch, 1, cond_dim)

        output
        ------
          y: tensor, (batch, length, in_dim)
        """
        y_tmp = z
        for l_flow in self.m_flows[::-1]:
            y_tmp = l_flow.reverse(y_tmp, cond)
        y = self.m_squeeze.reverse(y_tmp)
        return y

class Blow(torch_nn.Module):
    """Blow
    """
    def __init__(self, cond_dim, num_blocks, num_flows_inblock,
                 conv_dim_channel, conv_kernel_size):
        """Blow(cond_dim, num_blocks, num_flows_inblock,
                conv_dim_channel, conv_kernel_size)

        Args
        ----
          cond_dim:, int, conditional feature dim, (batch, length, cond_dim)
          num_blocks: int, number of WaveGlowBlocks
          num_flows_inblock: int, number of flow steps in one WaveGlowBlock
          conv_dim_channel: int, dim of convolution layers channels
          conv_kernel_size: int, kernel size of the convolution layers

        This model defines:

        cond (global) ----- -> | ------> | --------> |
                               v         v           v
        y --------------> BlowBlock1 -> BlowBlock2 -> ... -> z
        """
        super(Blow, self).__init__()

        # input is assumed to be waveform
        self.m_input_dim = 1

        # save the dimension for get_z_noises
        self.m_z_dim = 0

        # define blocks
        tmp_squeezed_in_dim = self.m_input_dim
        tmp_flow_blocks = []
        for i in range(num_blocks):
            tmp_flow_blocks.append(
                BlowBlock(
                    tmp_squeezed_in_dim, cond_dim, num_flows_inblock,
                    conv_dim_channel, conv_kernel_size))

            tmp_squeezed_in_dim = tmp_flow_blocks[-1].get_out_feat_dim()
        self.m_z_dim = tmp_squeezed_in_dim

        self.m_flowblocks = torch_nn.ModuleList(tmp_flow_blocks)

        # done
        return

    def get_expected_squeeze_length(self, wave_length):
        """length = get_expected_squeeze_length(self, wave_length)
        Return expected length of latent z

        input
        -----
          wave_length: int, length of original waveform

        output
        ------
          length: int, length of latent z
        """

        length = wave_length
        for glowblock in self.m_flowblocks:
            length = glowblock.get_expected_squeeze_length(length)
        return length

    def _normal_lh(self, noise):
        # likelihood of normal distribution on the given noise
        return -0.5 * np.log(2 * np.pi) - 0.5 * noise ** 2

    def forward(self, y, cond):
        """z, neg_logp_y, logp_z, logdet = Blow.forward(y, cond)

        cond (global) ----- -> | ------> | --------> |
                               v         v           v
        y --------------> BlowBlock1 -> BlowBlock2 -> ... -> z

        input
        -----
          y: tensor, (batch, waveform_length, 1)
          cond: tensor,  (batch, 1, cond_dim)

        output
        ------
          z: tensor
          neg_logp_y: scalar, - log p(y)
          logp_z: scalar, -log N(z), summed over one data sequence, but averaged
                  over batch.
          logdet: scalar, -|det dH(.)/dy|, summed over one data sequence,
                  but averaged
                  over batch.
        """

        # Rather than summing the likelihood and divide it by the number of
        #  data in the final step, we divide this factor from the likelihood
        #  caculating by each flow step and sum the scaled likelihood.
        # Two methods are equivalent, but the latter may prevent numerical
        #  overflow of the likelihood value for long sentences
        factor = np.prod([dim for dim in y.shape])

        # flows
        log_detjac = 0
        log_pz = 0

        x_tmp = y
        for m_block in self.m_flowblocks:
            x_tmp, log_detjac_tmp = m_block(
                x_tmp, cond, factor)

            # accumulate log det jacobian
            log_detjac += log_detjac_tmp

        z_tmp = x_tmp
        # compute N(z; 0, I)
        # accumulate log_N(z; 0, I) only if it is valid
        if z_tmp is not None:
            log_pz += nii_glow.sum_over_keep_batch2(
                self._normal_lh(z_tmp), factor)

        # average over batch and data points
        neg_logp_y = -(log_pz + log_detjac).sum()
        return z_tmp, neg_logp_y, \
            log_pz.sum(), log_detjac.sum()

    def reverse(self, z, cond):
        """y = Blow.reverse(z_bags, cond)

        cond (global) ----- -> | ------> | --------> |
                               v         v           v
        y <--------------- BlowBlock1 <- BlowBlock2 <- ... <- z

        input
        -----
          z: tensor, shape decided by the model configuration
          cond: tensor,  (batch, 1, cond_dim)

        output
        ------
          y: tensor, (batch, waveform_length, 1)
        """
        # initial
        y_tmp = z
        for m_block in self.m_flowblocks[::-1]:
            y_tmp = m_block.reverse(y_tmp, cond)
        return y_tmp

    def get_z_noises(self, length, noise_std=0.7, batchsize=1):
        """z_bags = Blow.get_z_noises(length, noise_std=0.7, batchsize=1)
        Return random noise for random sampling

        input
        -----
          length: int, length of target waveform (without squeeze)
          noise_std: float, std of Gaussian noise, default 0.7
          batchsize: int, batch size of this random data, default 1

        output
        ------
          z: tensor, shape decided by the network

        Blow.reverse(z, cond) can be used to generate waveform
        """
        squeeze_length = self.get_expected_squeeze_length(length)

        device = next(self.parameters()).device

        z_tmp = torch.randn(
            [batchsize, squeeze_length, self.m_z_dim],
            dtype=nii_io_conf.d_dtype,
            device=device)
        return z_tmp

if __name__ == "__main__":
    print("Definition of Blow")


r"""
Probabilistic / decoder modules for DIRECTi
"""

import typing
import abc

import numpy as np
import tensorflow as tf

from . import module, nn

class ProbModel(module.Module):
    r"""
    Abstract base class for generative model modules.
    """
    def __init__(
            self, h_dim: int = 128, depth: int = 1, dropout: float = 0.0,
            lambda_reg: float = 0.0, fine_tune: bool = False,
            deviation_reg: float = 0.0, name: str = "ProbModel"
    ) -> None:
        super(ProbModel, self).__init__(name=name)
        self.h_dim = h_dim
        self.depth = depth
        self.dropout = dropout
        self.lambda_reg = lambda_reg
        self.fine_tune = fine_tune
        self.deviation_reg = deviation_reg
        self.deviation_regularizer = \
            (lambda x: self.deviation_reg * tf.reduce_mean(tf.square(x))) \
            if self.fine_tune and self.deviation_reg > 0 else None

    @staticmethod
    def _normalize(  # pylint: disable=unused-argument
            x: typing.Union[np.ndarray, tf.Tensor],
            library_size: typing.Union[np.ndarray, tf.Tensor]
    ) -> typing.Union[np.ndarray, tf.Tensor]:  # pragma: no cover
        return x

    @staticmethod
    def _add_noise(  # pylint: disable=unused-argument
            x: typing.Union[np.ndarray, tf.Tensor],
            random_state: typing.Optional[np.random.RandomState] = None
    ) -> typing.Union[np.ndarray, tf.Tensor]:  # pragma: no cover
        return x

    @staticmethod
    def _preprocess(x: tf.Tensor) -> tf.Tensor:
        return x

    def _loss(
            self, ref: tf.Tensor, latent: tf.Tensor, training_flag: tf.Tensor,
            tail_concat: typing.Optional[typing.List[tf.Tensor]] = None,
            scope: str = "decoder"
    ) -> tf.Tensor:
        with tf.variable_scope(f"{scope}/{self.scope_safe_name}"):
            dropout = np.zeros(self.depth)
            dropout[1:] = self.dropout  # No dropout for first layer
            mlp_kwargs = dict(
                dropout=dropout.tolist(), dense_kwargs=dict(
                    deviation_regularizer=self.deviation_regularizer
                ), training_flag=training_flag
            )
            ptr = nn.mlp(latent, [self.h_dim] * self.depth, **mlp_kwargs)
            ptr = (ptr if isinstance(ptr, list) else [ptr]) + (tail_concat or [])
            self.log_likelihood = self._log_likelihood(ref, ptr)
            self.mean_log_likelihood = tf.reduce_mean(self.log_likelihood, axis=1)  # feature size invariant
            raw_loss = tf.negative(tf.reduce_mean(self.mean_log_likelihood), name="raw_loss")
            regularized_loss = tf.add(
                raw_loss, self.lambda_reg * self._build_regularizer(),
                name="regularized_loss"
            )
        self.vars_to_save += tf.get_collection(
            tf.GraphKeys.GLOBAL_VARIABLES, f"{scope}/{self.scope_safe_name}")
        tf.add_to_collection(tf.GraphKeys.LOSSES, raw_loss)
        tf.add_to_collection(tf.GraphKeys.LOSSES, regularized_loss)
        return regularized_loss

    @abc.abstractmethod
    def _log_likelihood(
            self, ref: tf.Tensor, pre_recon: typing.List[tf.Tensor]
    ) -> tf.Tensor:  # pragma: no cover
        raise NotImplementedError

    def _build_regularizer(self) -> tf.Tensor:
        return 0

    def _get_config(self) -> typing.Mapping:
        return {
            "h_dim": self.h_dim,
            "depth": self.depth,
            "dropout": self.dropout,
            "lambda_reg": self.lambda_reg,
            "fine_tune": self.fine_tune,
            "deviation_reg": self.deviation_reg,
            **super(ProbModel, self)._get_config()
        }

    def __bool__(self) -> bool:
        return True

class CountBased(ProbModel):  # pylint: disable=abstract-method

    @staticmethod
    def _normalize(
            x: typing.Union[np.ndarray, tf.Tensor],
            library_size: typing.Union[np.ndarray, tf.Tensor]
    ) -> typing.Union[tf.Tensor]:
        return x / (library_size / 10000)

    @staticmethod
    def _add_noise(
            x: typing.Union[np.ndarray, tf.Tensor],
            random_state: typing.Optional[np.random.RandomState] = None
    ) -> typing.Union[np.ndarray, tf.Tensor]:
        if random_state is None:
            return tf.squeeze(tf.random_poisson(x, [1]), axis=0)
        else:
            return random_state.poisson(x)

    @staticmethod
    def _preprocess(x: tf.Tensor) -> tf.Tensor:
        return tf.log1p(x)

class NB(CountBased):  # Negative binomial
    r"""
    Build a Negative Binomial generative module.

    Parameters
    ----------
    h_dim
        Dimensionality of the hidden layers in the decoder MLP.
    depth
        Number of hidden layers in the decoder MLP.
    dropout
        Dropout rate.
    fine_tune
        Whether the module is used in fine-tuning.
    lambda_reg
        Regularization strength for the generative model parameters.
        Here log-scale variance of the scale parameter
        is regularized to improve numerical stability.
    deviation_reg
        Regularization strength for the deviation from original model weights.
    name
        Name of the module.
    """
    def __init__(
            self, h_dim: int = 128, depth: int = 1, dropout: float = 0.0,
            lambda_reg: float = 0.0, fine_tune: bool = False,
            deviation_reg: float = 0.0, name: str = "NB"
    ) -> None:
        super(NB, self).__init__(
            h_dim, depth, dropout, lambda_reg,
            fine_tune, deviation_reg, name=name
        )

    def _log_likelihood(
            self, ref: tf.Tensor, pre_recon: typing.List[tf.Tensor]
    ) -> tf.Tensor:
        recon_dim = ref.get_shape().as_list()[1]
        self.softmax_mu = tf.nn.softmax(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="softmax_mu_dense"
        ), name="softmax_mu")
        self.log_theta = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="log_theta_dense"
        ), name="log_theta")
        self.recon = mu = \
            self.softmax_mu * tf.reduce_sum(ref, axis=1, keepdims=True)
        return self._log_nb_positive(ref, mu, self.log_theta)

    def _build_regularizer(self) -> tf.Tensor:
        with tf.name_scope("regularization"):
            return tf.nn.moments(self.log_theta, axes=[0, 1])[1]

    @staticmethod
    def _log_nb_positive(
            x: tf.Tensor, mu: tf.Tensor, log_theta: tf.Tensor,
            eps: float = 1e-8
    ) -> tf.Tensor:
        with tf.name_scope("log_nb_positive"):
            theta = tf.exp(log_theta)
            return theta * log_theta \
                - theta * tf.log(theta + mu + eps) \
                + x * tf.log(mu + eps) - x * tf.log(theta + mu + eps) \
                + tf.lgamma(x + theta) - tf.lgamma(theta) \
                - tf.lgamma(x + 1)

class ZINB(NB):  # Zero-inflated negative binomial
    r"""
    Build a Zero-Inflated Negative Binomial generative module.

    Parameters
    ----------
    h_dim
        Dimensionality of the hidden layers in the decoder MLP.
    depth
        Number of hidden layers in the decoder MLP.
    dropout
        Dropout rate.
    fine_tune
        Whether the module is used in fine-tuning.
    lambda_reg
        Regularization strength for the generative model parameters.
        Here log-scale variance of the scale parameter
        is regularized to improve numerical stability.
    deviation_reg
        Regularization strength for the deviation from original model weights.
    name
        Name of the module.
    """
    def __init__(
            self, h_dim: int = 128, depth: int = 1, dropout: float = 0.0,
            lambda_reg: float = 0.0, fine_tune: bool = False,
            deviation_reg: float = 0.0, name: str = "ZINB"
    ) -> None:
        super(ZINB, self).__init__(
            h_dim, depth, dropout, lambda_reg,
            fine_tune, deviation_reg, name=name
        )

    def _log_likelihood(
            self, ref: tf.Tensor, pre_recon: typing.List[tf.Tensor]
    ) -> tf.Tensor:
        recon_dim = ref.get_shape().as_list()[1]
        self.softmax_mu = tf.nn.softmax(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="softmax_mu_dense"
        ), name="softmax_mu")
        self.log_theta = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="log_theta_dense"
        ), name="log_theta")
        self.pi = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="dropout_logit_dense"
        ), name="dropout_logit")
        self.recon = mu = \
            self.softmax_mu * tf.reduce_sum(ref, axis=1, keepdims=True)
        self.dropout_rate = tf.sigmoid(self.pi)
        return self._log_zinb_positive(ref, mu, self.log_theta, self.pi)

    @staticmethod
    def _log_zinb_positive(
            x: tf.Tensor, mu: tf.Tensor, log_theta: tf.Tensor,
            pi: tf.Tensor, eps: float = 1e-8
    ) -> tf.Tensor:
        r"""
        From scVI
        """
        with tf.name_scope("log_zinb_positive"):
            theta = tf.exp(log_theta)
            with tf.name_scope("case_zero"):
                case_zero = tf.nn.softplus(
                    - pi + theta * log_theta -
                    theta * tf.log(theta + mu + eps)
                ) - tf.nn.softplus(- pi)
            with tf.name_scope("case_non_zero"):
                case_non_zero = - pi - tf.nn.softplus(- pi) \
                    + theta * log_theta \
                    - theta * tf.log(theta + mu + eps) \
                    + x * tf.log(mu + eps) - x * tf.log(theta + mu + eps) \
                    + tf.lgamma(x + theta) - tf.lgamma(theta) \
                    - tf.lgamma(x + 1)
            with tf.name_scope("mixture"):
                mask = tf.cast(tf.less(x, eps), tf.float32)
                res = tf.identity(
                    tf.multiply(mask, case_zero) +
                    tf.multiply(1 - mask, case_non_zero),
                    name="likelihood")
            return res

class LN(CountBased):
    r"""
    Build a Log Normal generative module.

    Parameters
    ----------
    h_dim
        Dimensionality of the hidden layers in the decoder MLP.
    depth
        Number of hidden layers in the decoder MLP.
    dropout
        Dropout rate.
    lambda_reg
        NOT USED.
    fine_tune
        Whether the module is used in fine-tuning.
    deviation_reg
        Regularization strength for the deviation from original model weights.
    name
        Name of the module.
    """
    def __init__(
            self, h_dim: int = 128, depth: int = 1, dropout: float = 0.0,
            lambda_reg: float = 0.0, fine_tune: bool = False,
            deviation_reg: float = 0.0, name: str = "LN"
    ) -> None:
        super(LN, self).__init__(
            h_dim, depth, dropout, lambda_reg,
            fine_tune, deviation_reg, name=name
        )

    def _log_likelihood(
            self, ref: tf.Tensor, pre_recon: typing.List[tf.Tensor]
    ) -> tf.Tensor:
        recon_dim = ref.get_shape().as_list()[1]
        self.mu = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="mu_dense"
        ), name="mu")
        self.log_var = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="log_var_dense"
        ), name="log_var")
        self.recon = tf.expm1(self.mu)
        return self._log_ln_positive(
            tf.log1p(ref), self.mu, self.log_var)

    @staticmethod
    def _log_ln_positive(
            x: tf.Tensor, mu: tf.Tensor, log_var: tf.Tensor
    ) -> tf.Tensor:
        with tf.name_scope("log_ln"):
            return - 0.5 * (
                tf.square(x - mu) / tf.exp(log_var)
                + tf.log(2 * np.pi) + log_var
            )

class ZILN(LN):
    r"""
    Build a Zero-Inflated Log Normal generative module.

    Parameters
    ----------
    h_dim
        Dimensionality of the hidden layers in the decoder MLP.
    depth
        Number of hidden layers in the decoder MLP.
    dropout
        Dropout rate.
    lambda_reg
        NOT USED.
    fine_tune
        Whether the module is used in fine-tuning.
    deviation_reg
        Regularization strength for the deviation from original model weights.
    name
        Name of the module.
    """
    def __init__(
            self, h_dim: int = 128, depth: int = 1, dropout: float = 0.0,
            lambda_reg: float = 0.0, fine_tune: bool = False,
            deviation_reg: float = 0.0, name: str = "ZILN"
    ) -> None:
        super(ZILN, self).__init__(
            h_dim, depth, dropout, lambda_reg,
            fine_tune, deviation_reg, name=name
        )

    def _log_likelihood(
            self, ref: tf.Tensor, pre_recon: typing.List[tf.Tensor]
    ) -> tf.Tensor:
        recon_dim = ref.get_shape().as_list()[1]
        self.mu = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="mu_dense"
        ), name="mu")
        self.log_var = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="log_var_dense"
        ), name="log_var")
        self.pi = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="dropout_logit_dense"
        ), name="dropout_logit")
        self.recon = tf.expm1(self.mu)
        self.dropout_rate = tf.sigmoid(self.pi)
        return self._log_ziln_positive(
            tf.log1p(ref), self.mu, self.log_var, self.pi)

    @staticmethod
    def _log_ziln_positive(
            x: tf.Tensor, mu: tf.Tensor, log_var: tf.Tensor,
            pi: tf.Tensor, eps: float = 1e-8
    ) -> tf.Tensor:
        with tf.name_scope("log_ziln"):
            with tf.name_scope("case_zero"):
                case_zero = - tf.nn.softplus(- pi)
            with tf.name_scope("case_non_zero"):
                case_non_zero = - pi - tf.nn.softplus(- pi) - 0.5 * (
                    tf.square(x - mu) / tf.exp(log_var)
                    + tf.log(2 * np.pi) + log_var
                )
            with tf.name_scope("mixture"):
                mask = tf.cast(tf.less(x, eps), tf.float32)
                res = tf.identity(
                    tf.multiply(mask, case_zero) +
                    tf.multiply(1 - mask, case_non_zero),
                    name="likelihood"
                )
            return res

class MSE(ProbModel):

    def __init__(
            self, h_dim: int = 128, depth: int = 1, dropout: float = 0.0,
            lambda_reg: float = 0.0, fine_tune: bool = False,
            deviation_reg: float = 0.0, name="MSE"
    ) -> None:
        super(MSE, self).__init__(
            h_dim, depth, dropout, lambda_reg,
            fine_tune, deviation_reg, name=name
        )

    def _log_likelihood(
            self, ref: tf.Tensor, pre_recon: typing.List[tf.Tensor]
    ) -> tf.Tensor:
        recon_dim = ref.get_shape().as_list()[1]
        self.mu = tf.identity(nn.dense(
            pre_recon, recon_dim,
            deviation_regularizer=self.deviation_regularizer,
            scope="mu_dense"
        ), name="mu")
        return tf.negative(tf.square(ref - self.mu))


import os
import numpy as np
import h5py

from ...utils.collections import dotdict

def loadDefaultParams(Cmat=None, Dmat=None, lookupTableFileName=None, seed=None):
    """Load default parameters for a network of aLN nodes.
    :param Cmat: Structural connectivity matrix (adjacency matrix) of coupling strengths, will be normalized to 1. If not given, then a single node simulation will be assumed, defaults to None
    :type Cmat: numpy.ndarray, optional
    :param Dmat: Fiber length matrix, will be used for computing the delay matrix together with the signal transmission speed parameter `signalV`, defaults to None
    :type Dmat: numpy.ndarray, optional
    :param lookUpTableFileName: Filename of lookup table with aln non-linear transfer functions and other precomputed quantities., defaults to aln-precalc/quantities_cascade.h
    :type lookUpTableFileName: str, optional
    :param seed: Seed for the random number generator, defaults to None
    :type seed: int, optional

    :return: A dictionary with the default parameters of the model
    :rtype: dict
    """

    params = dotdict({})

    # Todo: Model metadata
    # recently added for easier simulation of aln and brian in pypet
    params.model = "aln"
    params.name = "aln"
    params.description = "Adaptive linear-nonlinear model of exponential integrate-and-fire neurons"

    # runtime parameters
    params.dt = 0.1  # ms 0.1ms is reasonable
    params.duration = 2000  # Simulation duration (ms)
    np.random.seed(seed)  # seed for RNG of noise and ICs
    # set seed to 0, pypet will complain otherwise
    if seed is None:
        seed = 0
    params.seed = seed

    # options
    params.warn = 0  # warn if limits of lookup tables are exceeded
    params.dosc_version = 0  # if 0, use exponential fit to linear response function
    params.distr_delay = 0  # if 1, use distributed delays instead of fixed
    params.filter_sigma = 0  # if 1, filter sigmae/sigmai
    params.fast_interp = 1  # if 1, Interpolate the value from the look-up table instead of taking the closest value

    # ------------------------------------------------------------------------
    # global whole-brain network parameters
    # ------------------------------------------------------------------------

    if Cmat is None:
        params.N = 1
        params.Cmat = np.zeros((1, 1))
        params.lengthMat = np.zeros((1, 1))

    else:
        params.Cmat = Cmat.copy()  # coupling matrix
        np.fill_diagonal(params.Cmat, 0)  # no self connections
        params.N = len(params.Cmat)  # number of nodes
        params.lengthMat = Dmat  # delay matrix

    # Signal transmission speed in mm/ms
    params.signalV = 20.0

    # PSP current amplitude in (mV/ms) (or nA/[C]) for global coupling
    # connections between areas
    params.c_gl = 0.3
    # number of incoming E connections (to E population) from each area
    params.Ke_gl = 250.0

    # ------------------------------------------------------------------------
    # local E-I node parameters
    # ------------------------------------------------------------------------

    # external input parameters:
    params.tau_ou = 5.0  # ms timescale of ornstein-uhlenbeck (OU) noise
    params.sigma_ou = 0.0  # mV/ms/sqrt(ms) intensity of OU oise
    params.mue_ext_mean = 0.4  # mV/ms mean external input current to E
    params.mui_ext_mean = 0.3  # mV/ms mean external input current to I

    # Ornstein-Uhlenbeck noise state variables, set to mean input
    # mue_ou will fluctuate around mue_ext_mean (mean of the OU process)
    params.mue_ou = params.mue_ext_mean * np.ones((params.N,))  # np.zeros((params.N,))
    params.mui_ou = params.mui_ext_mean * np.ones((params.N,))  # np.zeros((params.N,))

    # external neuronal firing rate input
    params.ext_exc_rate = 0.0  # kHz external excitatory rate drive
    params.ext_inh_rate = 0.0  # kHz external inhibiroty rate drive

    # externaln input currents, same as mue_ext_mean but can be time-dependent!
    params.ext_exc_current = 0.0  # external excitatory input current [mV/ms], C*[]V/s=[]nA
    params.ext_inh_current = 0.0  # external inhibiroty input current [mV/ms]

    # Fokker Planck noise (for N->inf)
    params.sigmae_ext = 1.5  # mV/sqrt(ms) (fixed, for now) [1-5] (Internal noise due to random coupling)
    params.sigmai_ext = 1.5  # mV/sqrt(ms) (fixed, for now) [1-5]

    # recurrent coupling parameters
    params.Ke = 800.0  # Number of excitatory inputs per neuron
    params.Ki = 200.0  # Number of inhibitory inputs per neuron

    # synaptic delays
    params.de = 4.0  # ms local constant delay "EE = IE"
    params.di = 2.0  # ms local constant delay "EI = II"

    # synaptic time constants
    params.tau_se = 2.0  # ms  "EE = IE", for fixed delays
    params.tau_si = 5.0  # ms  "EI = II"

    # time constant for distributed delays (untested)
    params.tau_de = 1.0  # ms  "EE = IE"
    params.tau_di = 1.0  # ms  "EI = II"

    # PSC amplitudes
    params.cee = 0.3  # mV/ms
    params.cie = 0.3  # AMPA
    params.cei = 0.5  # GABA BrunelWang2003
    params.cii = 0.5

    # Coupling strengths used in Cakan2020
    params.Jee_max = 2.43  # mV/ms
    params.Jie_max = 2.60  # mV/ms
    params.Jei_max = -3.3  # mV/ms [0-(-10)]
    params.Jii_max = -1.64  # mV/ms

    # neuron model parameters
    params.a = 0.0  # nS, can be 15.0
    params.b = 0.0  # pA, can be 40.0
    params.EA = -80.0  # mV
    params.tauA = 200.0  # ms

    # single neuron paramters - if these are changed, new transfer functions must be precomputed!
    params.C = 200.0  # pF
    params.gL = 10.0  # nS
    params.EL = -65.0  # mV
    params.DeltaT = 1.5  # mV
    params.VT = -50.0  # mV
    params.Vr = -70.0  # mV
    params.Vs = -40.0  # mV
    params.Tref = 1.5  # ms

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

    # Generate and set random initial conditions
    (
        mufe_init,
        mufi_init,
        IA_init,
        seem_init,
        seim_init,
        seev_init,
        seiv_init,
        siim_init,
        siem_init,
        siiv_init,
        siev_init,
        rates_exc_init,
        rates_inh_init,
    ) = generateRandomICs(params.N, seed)

    params.mufe_init = mufe_init  # (linear) filtered mean input
    params.mufi_init = mufi_init  #
    params.IA_init = IA_init  # adaptation current
    params.seem_init = seem_init  # mean of fraction of active synapses [0-1] (post-synaptic variable), chap. 4.2
    params.seim_init = seim_init  #
    params.seev_init = seev_init  # variance of fraction of active synapses [0-1]
    params.seiv_init = seiv_init  #
    params.siim_init = siim_init  #
    params.siem_init = siem_init  #
    params.siiv_init = siiv_init  #
    params.siev_init = siev_init  #
    params.rates_exc_init = rates_exc_init  #
    params.rates_inh_init = rates_inh_init  #

    # load precomputed aLN transfer functions from hdfs
    if lookupTableFileName is None:
        lookupTableFileName = os.path.join(os.path.dirname(__file__), "aln-precalc", "quantities_cascade.h5")

    hf = h5py.File(lookupTableFileName, "r")
    params.Irange = hf.get("mu_vals")[()]
    params.sigmarange = hf.get("sigma_vals")[()]
    params.dI = params.Irange[1] - params.Irange[0]
    params.ds = params.sigmarange[1] - params.sigmarange[0]

    params.precalc_r = hf.get("r_ss")[()][()]
    params.precalc_V = hf.get("V_mean_ss")[()]
    params.precalc_tau_mu = hf.get("tau_mu_exp")[()]
    params.precalc_tau_sigma = hf.get("tau_sigma_exp")[()]

    return params

def computeDelayMatrix(lengthMat, signalV, segmentLength=1):
    """
    Compute the delay matrix from the fiber length matrix and the signal
    velocity

        :param lengthMat:       A matrix containing the connection length in
            segment
        :param signalV:         Signal velocity in m/s
        :param segmentLength:   Length of a single segment in mm

        :returns:    A matrix of connexion delay in ms
    """

    normalizedLenMat = lengthMat * segmentLength
    if signalV > 0:
        Dmat = normalizedLenMat / signalV  # Interareal delays in ms
    else:
        Dmat = lengthMat * 0.0
    return Dmat

def generateRandomICs(N, seed=None):
    """ Generates random Initial Conditions for the interareal network

        :params N:  Number of area in the large scale network

        :returns:   A tuple of 9 N-length numpy arrays representining:
                        mufe_init, IA_init, mufi_init, sem_init, sev_init,
                        sim_init, siv_init, rates_exc_init, rates_inh_init
    """
    np.random.seed(seed)

    mufe_init = 3 * np.random.uniform(0, 1, (N,))  # mV/ms
    mufi_init = 3 * np.random.uniform(0, 1, (N,))  # mV/ms
    IA_init = 200.0 * np.random.uniform(0, 1, (N,))  # pA
    seem_init = 0.5 * np.random.uniform(0, 1, (N,))
    seim_init = 0.5 * np.random.uniform(0, 1, (N,))
    seev_init = 0.001 * np.random.uniform(0, 1, (N,))
    seiv_init = 0.001 * np.random.uniform(0, 1, (N,))
    siim_init = 0.5 * np.random.uniform(0, 1, (N,))
    siem_init = 0.5 * np.random.uniform(0, 1, (N,))
    siiv_init = 0.01 * np.random.uniform(0, 1, (N,))
    siev_init = 0.01 * np.random.uniform(0, 1, (N,))
    rates_exc_init = 0.01 * np.random.uniform(0, 1, (N, 1))
    rates_inh_init = 0.01 * np.random.uniform(0, 1, (N, 1))

    return (
        mufe_init,
        mufi_init,
        IA_init,
        seem_init,
        seim_init,
        seev_init,
        seiv_init,
        siim_init,
        siem_init,
        siiv_init,
        siev_init,
        rates_exc_init,
        rates_inh_init,
    )


import struct
import numpy as np
from astropy import units as u
from astropy.coordinates import Angle
import os

###
# Header parsing
###

# Dictionary of allowed keywords and their types
# Here are the keywordss that a filter bank file may
# contain.  Items marked with "[*]" are not yet # supported.  See docs for
# indivisuabl attribtues for more detailed info.
#
#   * telescope_id (int): 0=fake data; 1=Arecibo; 2=Ooty... others to be added
#   * machine_id (int): 0=FAKE; 1=PSPM; 2=WAPP; 3=OOTY... others to be added
#   * data_type (int): 1=blimpy; 2=time series... others to be added
#   * rawdatafile (string): the name of the original data file
#   * source_name (string): the name of the source being observed by the telescope
#   * barycentric (int): equals 1 if data are barycentric or 0 otherwise
#   * pulsarcentric (int): equals 1 if data are pulsarcentric or 0 otherwise
#   * az_start (double): telescope azimuth at start of scan (degrees)
#   * za_start (double): telescope zenith angle at start of scan (degrees)
#   * src_raj (double): right ascension (J2000) of source (hours, converted from hhmmss.s)
#   * src_dej (double): declination (J2000) of source (degrees, converted from ddmmss.s)
#   * tstart (double): time stamp (MJD) of first sample
#   * tsamp (double): time interval between samples (s)
#   * nbits (int): number of bits per time sample
#   * nsamples (int): number of time samples in the data file (rarely used any more)
#   * fch1 (double): centre frequency (MHz) of first blimpy channel
#   * foff (double): blimpy channel bandwidth (MHz)
#   * FREQUENCY_START [*] (character): start of frequency table (see below for explanation)
#   * fchannel [*] (double): frequency channel value (MHz)
#   * FREQUENCY_END [*] (character): end of frequency table (see below for explanation)
#   * nchans (int): number of blimpy channels
#   * nifs (int): number of seperate IF channels
#   * refdm (double): reference dispersion measure (pc/cm**3)
#   * period (double): folding period (s)
#   * nbeams (int):total number of beams (?)
#   * ibeam (int): number of the beam in this file (?)

header_keyword_types = {
    'telescope_id' : '<l',
    'machine_id'   : '<l',
    'data_type'    : '<l',
    'barycentric'  : '<l',
    'pulsarcentric': '<l',
    'nbits'        : '<l',
    'nsamples'     : '<l',
    'nchans'       : '<l',
    'nifs'         : '<l',
    'nbeams'       : '<l',
    'ibeam'        : '<l',
    'rawdatafile'  : 'str',
    'source_name'  : 'str',
    'az_start'     : '<d',
    'za_start'     : '<d',
    'tstart'       : '<d',
    'tsamp'        : '<d',
    'fch1'         : '<d',
    'foff'         : '<d',
    'refdm'        : '<d',
    'period'       : '<d',
    'src_raj'      : 'angle',
    'src_dej'      : 'angle',
}

def len_header(filename):
    """ Return the length of the blimpy header, in bytes

    Args:
        filename (str): name of file to open

    Returns:
        idx_end (int): length of header, in bytes
    """
    with open(filename, 'rb') as f:
        header_sub_count = 0
        eoh_found = False
        while not eoh_found:
            header_sub = f.read(512)
            header_sub_count += 1
            if b'HEADER_END' in header_sub:
                idx_end = header_sub.index(b'HEADER_END') + len(b'HEADER_END')
                eoh_found = True
                break

        idx_end = (header_sub_count -1) * 512 + idx_end
    return idx_end

def read_next_header_keyword(fh):
    """

    Args:
        fh (file): file handler

    Returns:
    """
    n_bytes = np.frombuffer(fh.read(4), dtype='uint32')[0]

    if n_bytes > 255:
        n_bytes = 16

    keyword = fh.read(n_bytes).decode('ascii')

    if keyword in ('HEADER_START', 'HEADER_END'):
        return keyword, 0, fh.tell()
    dtype = header_keyword_types[keyword]
    idx = fh.tell()
    if dtype == '<l':
        val = struct.unpack(dtype, fh.read(4))[0]
    if dtype == '<d':
        val = struct.unpack(dtype, fh.read(8))[0]
    if dtype == 'str':
        str_len = np.frombuffer(fh.read(4), dtype='uint32')[0]
        val = fh.read(str_len).decode('ascii')
    if dtype == 'angle':
        val = struct.unpack('<d', fh.read(8))[0]
        val = fil_double_to_angle(val)
        if keyword == 'src_raj':
            val = Angle(val, unit=u.hour)
        else:
            val = Angle(val, unit=u.deg)
    return keyword, val, idx

def is_filterbank(filename):
    """ Open file and confirm if it is a filterbank file or not. """
    with open(filename, 'rb') as fh:
        is_fil = True

        # Check this is a blimpy file
        try:
            keyword, value, idx = read_next_header_keyword(fh)
            try:
                assert keyword == 'HEADER_START'
            except AssertionError:
                is_fil = False
        except KeyError:
            is_fil = False
        return is_fil

def read_header(filename, return_idxs=False):
    """ Read blimpy header and return a Python dictionary of key:value pairs

    Args:
        filename (str): name of file to open

    Optional args:
        return_idxs (bool): Default False. If true, returns the file offset indexes
                            for values

    returns

    """
    with open(filename, 'rb') as fh:
        header_dict = {}
        header_idxs = {}

        # Check this is a blimpy file
        keyword, value, idx = read_next_header_keyword(fh)

        try:
            assert keyword == 'HEADER_START'
        except AssertionError:
            raise RuntimeError("Not a valid blimpy file.")

        while True:
            keyword, value, idx = read_next_header_keyword(fh)
            if keyword == 'HEADER_END':
                break
            header_dict[keyword] = value
            header_idxs[keyword] = idx

    if return_idxs:
        return header_idxs
    return header_dict

def fix_header(filename, keyword, new_value):
    """ Apply a quick patch-up to a Filterbank header by overwriting a header value

    Args:
        filename (str): name of file to open and fix. WILL BE MODIFIED.
        keyword (stt):  header keyword to update
        new_value (long, double, angle or string): New value to write.

    Notes:
        This will overwrite the current value of the blimpy with a desired
        'fixed' version. Note that this has limited support for patching
        string-type values - if the length of the string changes, all hell will
        break loose.

    """

    # Read header data and return indexes of data offsets in file
    hd = read_header(filename)
    hi = read_header(filename, return_idxs=True)
    idx = hi[keyword]

    # Find out the datatype for the given keyword
    dtype = header_keyword_types[keyword]
    dtype_to_type = {'<l'  : np.int32,
                     'str' : bytes,
                     '<d'  : np.float64,
                     'angle' : to_sigproc_angle}
    value_dtype = dtype_to_type[dtype]

    # Generate the new string
    if isinstance(value_dtype, bytes):
        if len(hd[keyword]) == len(new_value):
            val_str = np.int32(len(new_value)).tostring() + new_value
        else:
            raise RuntimeError("String size mismatch. Cannot update without rewriting entire file.")
    else:
        val_str = value_dtype(new_value).tostring()

    # Write the new string to file
    with open(filename, 'rb+') as fh:
        fh.seek(idx)
        fh.write(val_str)

def fil_double_to_angle(angle):
    """ Reads a little-endian double in ddmmss.s (or hhmmss.s) format and then
    converts to Float degrees (or hours).  This is primarily used to read
    src_raj and src_dej header values. """

    negative = (angle < 0.0)
    angle = np.abs(angle)

    dd = np.floor((angle / 10000))
    angle -= 10000 * dd
    mm = np.floor((angle / 100))
    ss = angle - 100 * mm
    dd += mm/60.0 + ss/3600.0

    if negative:
        dd *= -1

    return dd

###
# sigproc writing functions
###

def to_sigproc_keyword(keyword, value=None):
    """ Generate a serialized string for a sigproc keyword:value pair

    If value=None, just the keyword will be written with no payload.
    Data type is inferred by keyword name (via a lookup table)

    Args:
        keyword (str): Keyword to write
        value (None, float, str, double or angle): value to write to file

    Returns:
        value_str (str): serialized string to write to file.
    """
    if value is None:
        return np.int32(len(keyword)).tobytes() + keyword.encode('ascii')
    dtype = header_keyword_types[keyword]

    dtype_to_type = {'<l'  : np.int32,
                     'str' : str,
                     '<d'  : np.float64,
                     'angle' : to_sigproc_angle}

    value_dtype = dtype_to_type[dtype]

    if isinstance(value, str):
        value = value.encode('ascii')
    if value_dtype is str:
        return np.int32(len(keyword)).tobytes() + keyword.encode('ascii') + np.int32(len(value)).tobytes() + value
    return np.int32(len(keyword)).tobytes() + keyword.encode('ascii') + value_dtype(value).tobytes()

def generate_sigproc_header(f):
    """ Generate a serialzed sigproc header which can be written to disk.

    Args:
        f (Filterbank object): Filterbank object for which to generate header

    Returns:
        header_str (str): Serialized string corresponding to header
    """

    header_string = b''
    header_string += to_sigproc_keyword('HEADER_START')

    for keyword in f.header.keys():
        if keyword == 'src_raj':
            header_string += to_sigproc_keyword('src_raj')  + to_sigproc_angle(f.header['src_raj'])
        elif keyword == 'src_dej':
            header_string += to_sigproc_keyword('src_dej')  + to_sigproc_angle(f.header['src_dej'])
        elif keyword in ('az_start', 'za_start'):
            header_string += to_sigproc_keyword(keyword)  + np.float64(f.header[keyword]).tobytes()
        elif keyword not in header_keyword_types.keys():
            pass
        else:
            header_string += to_sigproc_keyword(keyword, f.header[keyword])

    header_string += to_sigproc_keyword('HEADER_END')
    return header_string

def to_sigproc_angle(angle_val):
    """ Convert an astropy.Angle to the ridiculous sigproc angle format string. """
    x = str(angle_val)

    if '.' in x:
        if 'h' in x:
            d, m, s, ss = int(x[0:x.index('h')]), int(x[x.index('h')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('.')]), float(x[x.index('.'):x.index('s')])
        if 'd' in x:
            d, m, s, ss = int(x[0:x.index('d')]), int(x[x.index('d')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('.')]), float(x[x.index('.'):x.index('s')])
    else:
        if 'h' in x:
            d, m, s = int(x[0:x.index('h')]), int(x[x.index('h')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('s')])
        if 'd' in x:
            d, m, s = int(x[0:x.index('d')]), int(x[x.index('d')+1:x.index('m')]), \
            int(x[x.index('m')+1:x.index('s')])
        ss = 0
    num = str(d).zfill(2) + str(m).zfill(2) + str(s).zfill(2)+ '.' + str(ss).split(".")[-1]
    return np.float64(num).tobytes()

def calc_n_ints_in_file(filename):
    """ Calculate number of integrations in a given file """
    # Load binary data
    h = read_header(filename)
    n_chans = h['nchans']
    n_ifs   = h['nifs']
    idx_data = len_header(filename)
    filesize = os.path.getsize(filename)
    n_bytes_data = filesize - idx_data

    if h['nbits'] == 2:
        n_ints = int(4 * n_bytes_data / (n_chans * n_ifs))
    elif h['nbits'] == 4:
        n_ints = int(2 * n_bytes_data / (n_chans * n_ifs))
    else:
        n_bytes  = int(h['nbits'] / 8)
        n_ints = int(n_bytes_data / (n_bytes * n_chans * n_ifs))

    return n_ints


from math import ceil

import blosc
from blosc import compress_ptr, decompress_ptr, get_cbuffer_sizes
from numpy import ndarray

def compress_array(array: ndarray, clevel: int = 3, compressor: str = "lz4", min_num_chunks: int = 0) -> bytes:
    """
    Compresses an arbitrary ndarray that supports the '__array_interface__' into a Blosc compressed buffer.

    Parameters
    ----------
    array: Array to compress
    clevel: Compression level
    compressor: compressor (any supported by Blosc)
    min_num_chunks: Minimum number of chunks to split array into before compression.

    Returns
    -------
    Compressed buffer (as bytes)

    """
    array_address = array.__array_interface__["data"][0]
    array_length = array.nbytes
    end_address = array_address + array_length
    itemsize = array.itemsize

    # Blosc does not support 64bit buffers, only 32 bit buffers:
    max_chunk_length = blosc.MAX_BUFFERSIZE

    # Let's avoid small trailing chunks and make them all about teh same size:
    number_of_chunks = max(2, ceil(array_length / max_chunk_length)) if array_length > max_chunk_length else 1
    number_of_chunks = max(min_num_chunks, number_of_chunks)
    approx_chunk_length_in_bytes = array_length // number_of_chunks

    # Let's make sure that the chunks are aligned to the data type:
    approx_chunk_length_in_bytes = (approx_chunk_length_in_bytes // itemsize) * itemsize

    # Let' make the chunks large enough that we are garanteed to cover the whole array:
    approx_chunk_length_in_bytes += number_of_chunks * itemsize

    # this will hold the compressed data:
    compressed_buffer = bytearray()

    # we start reading the array data here:
    chunk_address = array_address

    # We loop through each chunk:
    while chunk_address < end_address:
        chunk_length_in_bytes = min(approx_chunk_length_in_bytes, end_address - chunk_address)
        compressed_chunk = compress_ptr(
            address=chunk_address,
            items=chunk_length_in_bytes // itemsize,
            typesize=itemsize,
            clevel=clevel,
            shuffle=blosc.BITSHUFFLE,
            cname=compressor,
        )
        if number_of_chunks == 1:
            # If there is only one chunk, let's not be complicated about it:
            compressed_buffer = compressed_chunk
        else:
            # If not, so bit it:
            compressed_buffer.extend(compressed_chunk)
        chunk_address += chunk_length_in_bytes

    return bytes(compressed_buffer)

def decompress_array(compressed_bytes, out_array: ndarray = None) -> ndarray:
    """
    Decompresses an array compressed with the 'compress_array' function.

    Parameters
    ----------
    compressed_bytes: buffer containing compressed data
    out_array: array of _correct_ size to put the decompressed daat in.

    Returns
    -------
    Same array as passed as 'out_array'

    """

    # get the array pointer and length in bytes:
    array_address = out_array.__array_interface__["data"][0]

    # prepare the basics:
    num_of_compressed_bytes = len(compressed_bytes)
    offset_compressed = 0
    address_decompressed = array_address

    while num_of_compressed_bytes - offset_compressed > 32:

        # This is the BLOSC header:
        blosc_header = bytes(compressed_bytes[offset_compressed : offset_compressed + 32])

        # we check how large is the first chunk available:
        num_decompressed_bytes, num_compressed_bytes, _ = get_cbuffer_sizes(blosc_header)

        # we prepare the corresponding region of the array:
        compressed_chunk = compressed_bytes[offset_compressed : offset_compressed + num_compressed_bytes]

        # do the actual decompression:
        decompress_ptr(compressed_chunk, address_decompressed)

        # we move the pointers forward:
        offset_compressed += num_compressed_bytes
        address_decompressed += num_decompressed_bytes

    return out_array


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ----------------------------------------------------------------------
# Copyright 2017-2020 Airinnova AB and the PyTornado authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------

# Authors:
# * Alessandro Gastaldi
# * Aaron Dettmann

"""
Functions for the discretisation of the aircraft geometry into panels.

Developed for AIRINNOVA AB, Stockholm, Sweden.
"""

import logging
from math import ceil

import numpy as np
import scipy.linalg.lapack as lapack
from commonlibs.math.vectors import axis_rot_matrix
from commonlibs.math.interpolation import lin_interpol

from pytornado.objects.vlm_struct import VLMLattice
import pytornado.aero.c_vlm as c_vlm
import pytornado.objects.objecttools as ot
from pytornado.objects.vlm_struct import BookKeepingEntry
from pytornado.objects.aircraft import get_abs_segment_point_coords

logger = logging.getLogger(__name__)

MIN_AUTOPANELS = 1

def set_autopanels(aircraft, settings):
    """
    Automatically set chord- and spanwise discretisation settings

    Args:
        :aircraft: (object) data structure for aircraft geometry
        :autopanels_c: (int) number of chordwise panels on the main wing
        :autopanels_s: (int) number of spanwise panels on the main wing
    """

    autopanels_c = settings.settings.get('vlm_autopanels_c', MIN_AUTOPANELS)
    autopanels_s = settings.settings.get('vlm_autopanels_s', MIN_AUTOPANELS)

    for this_segment, _ in ot.all_segments(aircraft):
        segment = this_segment[2]

        if segment.panels['num_c'] is None:
            segment.panels['num_c'] = autopanels_c

        if segment.panels['num_s'] is None:
            wing_span = segment.parent_wing.span
            segment_span = segment.geometry['span']
            segment.panels['num_s'] = ceil((segment_span/wing_span)*autopanels_s)

    for this_control, this_wing in ot.all_controls(aircraft):
        control = this_control[2]

        if control.panels['num_c'] is None:
            control.panels['num_c'] = autopanels_c

def pre_panelling(aircraft):
    """
    Create subdivisions and subareas for all aircraft wings

    Note:
        * This routine divides the wing into subdivisions and subareas
          in order to generate a suitable mesh for wing with control surfaces.
        * In a first step "mandatory" subdivisions are made: The wing is divided
          into a minimum amount of subareas according to the leading and trailing
          edge control surfaces.
        * In a second step further spanwise subdivisions are added.

    Args:
        :aircraft: (obj) aircraft object
    """

    # TODO:
    # - Potential problem:
    #   * The algorithm is based on "correct ordering" of segments:
    #     Segments must be ordered from root to tip (check if this always given!?)

    # ===== PART ONE (MANDATORY SUBDIVISIONS) =====

    # For each control we must add suitable subdivisions
    for this_control, this_wing in ot.all_controls(aircraft):
        control = this_control[2]
        wing = this_wing[2]

        # Segment names on which control surface edges are located
        segment_inner_name = control.segment_uid['inner']
        segment_outer_name = control.segment_uid['outer']

        # Control surface geometry
        eta_inner = control.rel_vertices['eta_inner']
        eta_outer = control.rel_vertices['eta_outer']
        xsi_inner = control.rel_vertices['xsi_inner']
        xsi_outer = control.rel_vertices['xsi_outer']

        # Hinge axis
        xsi_h1 = control.rel_hinge_vertices['xsi_inner']
        xsi_h2 = control.rel_hinge_vertices['xsi_outer']

        # ----- CASE (A) -----
        # The left and right edge of the control are located on SAME segment
        if segment_inner_name == segment_outer_name:
            wing.segments[segment_inner_name].add_subdivision_for_control(
                    eta_inner, eta_outer, control,
                    xsi_inner, xsi_outer, xsi_h1, xsi_h2)

        # ----- CASE (B) -----
        # - The control surface spans over one or more segment borders
        # - Now we will make a list of segments over which the control spans
        # - We will interpolate the position of the control and the hinge
        #   axis at segment borders
        else:
            # Create a list of segments which contain the control
            # list_of_segments[0]: segment_uid
            # list_of_segments[1]: eta_inner
            # list_of_segments[2]: eta_outer
            # list_of_segments[3]: xsi_inner
            # list_of_segments[4]: xsi_outer
            # list_of_segments[5]: xsi_h1
            # list_of_segments[6]: xsi_h2
            list_of_segments = []

            # Flags to indicate that the inner or outer control positions have been set
            inner_set = False
            outer_set = False

            # To start with we use averaged values for xsi (geometry and hinge axis)
            xsi_avg = (xsi_inner + xsi_outer)/2
            xsi_h_avg = (xsi_h1 + xsi_h2)/2

            for segment_uid in wing.segments.keys():
                if segment_uid == segment_inner_name:
                    inner_set = True
                    # Note: eta_outer = 1
                    list_of_segments.append([segment_uid, eta_inner, 1, xsi_inner, xsi_avg, xsi_h1, xsi_h_avg])
                    continue

                elif inner_set and not outer_set:
                    # Note: eta_inner = 0
                    # Note: eta_outer = 1
                    list_of_segments.append([segment_uid, 0, 1, xsi_avg, xsi_avg, xsi_h_avg, xsi_h_avg])

                    # If we are on the last segment we must update some incorrectly set values
                    if segment_uid == segment_outer_name:
                        outer_set = True
                        list_of_segments[-1][2] = eta_outer
                        list_of_segments[-1][4] = xsi_outer
                        list_of_segments[-1][6] = xsi_h2
                        break

            # Potentially, we must readjust the control surface geometry/hinge axis at borders
            if (xsi_inner != xsi_outer) or (xsi_h1 != xsi_h2):
                # Let's first compute the "length" of the control surface
                control_len = [0, ]

                for row in list_of_segments:
                    segment_uid, eta_i, eta_o, xsi_i, xsi_o, xsi_hi, xsi_ho = row

                    segment = wing.segments[segment_uid]
                    segment_vertices = segment.vertices

                    a = get_abs_segment_point_coords(segment_vertices, eta_i, xsi_i)
                    b = get_abs_segment_point_coords(segment_vertices, eta_o, xsi_o)
                    ab = b - a

                    # l: total length of control
                    l = control_len[-1]
                    control_len.append(l + np.sqrt(np.dot(ab, ab)))

                l = control_len[-1]

                # Now, we update the xsi values using linear interpolation
                for i, row in enumerate(list_of_segments):
                    segment_uid, eta_i, eta_o, xsi_i, xsi_o, xsi_hi, xsi_ho = row

                    # Update the xsi values
                    l_i = control_len[i]
                    l_o = control_len[i+1]

                    xsi_i = lin_interpol((xsi_inner, xsi_outer), (0, l), l_i)
                    xsi_o = lin_interpol((xsi_inner, xsi_outer), (0, l), l_o)
                    xsi_hi = lin_interpol((xsi_h1, xsi_h2), (0, l), l_i)
                    xsi_ho = lin_interpol((xsi_h1, xsi_h2), (0, l), l_o)

                    list_of_segments[i] = [segment_uid, eta_i, eta_o, xsi_i, xsi_o, xsi_hi, xsi_ho]

            # Finally, we create the subdivisions using our list
            for row in list_of_segments:
                segment_uid, eta_i, eta_o, xsi_i, xsi_o, xsi_h1, xsi_h2 = row

                wing.segments[segment_uid].add_subdivision_for_control(
                        eta_i, eta_o, control, xsi_i, xsi_o, xsi_h1, xsi_h2)

    # ===== PART TWO (ADDITIONAL SPANWISE SUBDIVISIONS) =====

    # - Adding additional spanwise subdivisions is done here in Python rather
    #   than in the C code as it is more convenient to keep track of which
    #   parts (subareas) of the discretised surface have which functions
    for this_segment, _ in ot.all_segments(aircraft):
        segment = this_segment[2]

        for eta in np.linspace(0, 1, segment.panels['num_s']+1):
            if (eta == 0) or (eta == 1):
                continue
            segment.add_subdivision(eta, eta, ignore_inval_eta=True)

def gen_lattice(aircraft, state, settings, make_new_subareas=True):
    """
    Generate aircraft lattice

    Perform count of number of wings, segments, controls, panels and strips.
    Pre-allocate memory for lattice data, which is directly operated on in C.

    The function `py2c_lattice` is called which generates the VLM lattice.
    The following lattice data (for all panel) is generated:

        * :lattice.p: panel corner points
        * :lattice.v: panel vortex filament endpoints
        * :lattice.c: panel collocation point
        * :lattice.n: panel normal vector
        * :lattice.a: panel surface area

    When `py2c_lattice` is called it takes the following input arguments:

        * :lattice: pre-allocated memory for the struct described above
        * :array_segments: segment corner points (N*4*3)
        * :array_airfoils: file names for airfoils at inner and outer segment (N*2)
        * :array_symmetry: segment symmetry information (N)
        * :array_panels: number of chordwise and spanwise panels for each segment (N*2)

    Display lattice metrics in console and log file.

    Args:
        :aircraft: (object) data structure for aircraft geometry
        :state: (object) data structure for flight state
        :settings: (object) data structure for execution settings
        :make_new_subareas: Flag

    Returns:
        :lattice: (object) data structure for VLM lattice
    """

    if make_new_subareas:
        pre_panelling(aircraft)

    # Start the panel bookkeping with a clean slate
    lattice = VLMLattice()
    lattice.clean_bookkeeping()

    logger.info("Getting lattice information ... ")

    num_subareas = 0
    num_r = 0  # total number of panel strips
    num_p = 0  # total number of panels

    # PANEL COUNT AND BOOK KEEPING
    for this_subarea, _, this_segment, this_wing in ot.all_subareas(aircraft):
        wing = this_wing[2]
        segment = this_segment[2]
        subarea = this_subarea[2]

        num_subareas += 1

        pan_idx1 = num_p

        # A subarea only has chordwise subdivisions
        num_chordwise_panels = subarea.parent_control.panels['num_c'] if \
                               subarea.parent_control is not None else \
                               segment.panels['num_c']

        # If a subdivisions has a flap and/or slat we can reduce the number of
        # chordwise subdivision
        if subarea.type == 'segment':
            num_chordwise_panels = ceil(subarea.rel_length*num_chordwise_panels)

        # NOTE: now num_r and num_p are the same (improve !?)
        num_r += num_chordwise_panels
        num_p += num_chordwise_panels

        lattice.update_bookkeeping(
                BookKeepingEntry(subarea, range(pan_idx1, num_p), num_chordwise_panels, mirror=False))

        if wing.symmetry:
            num_subareas += 1

            pan_idx1 = num_p
            num_r += num_chordwise_panels
            num_p += num_chordwise_panels

            lattice.update_bookkeeping(
                    BookKeepingEntry(subarea, range(pan_idx1, num_p), num_chordwise_panels, mirror=True))

    # Make sure integers are stored as integers
    num_wings = int(ot.count_all_wings(aircraft))
    num_controls = int(ot.count_all_controls(aircraft))
    num_subareas = int(num_subareas)
    num_r = int(num_r)
    num_p = int(num_p)

    lattice.info['num_wings'] = num_wings
    lattice.info['num_segments'] = num_subareas
    lattice.info['num_controls'] = num_controls
    lattice.info['num_strips'] = num_r
    lattice.info['num_panels'] = num_p

    logger.info("Pre-allocating lattice memory...")

    array_subareas = np.zeros((num_subareas, 4, 3), dtype=float, order='C')
    array_symmetry = np.zeros((num_subareas), dtype=int, order='C')
    array_panels = np.ones((num_subareas, 2), dtype=int, order='C')

    i = 0
    for entry in lattice.panel_bookkeeping:
        subarea = entry.subarea
        pan_idx = entry.pan_idx
        mirror = entry.mirror

        vertices = subarea.abs_vertices(mirror)

        array_subareas[i, 0, :] = vertices['a']
        array_subareas[i, 1, :] = vertices['b']
        array_subareas[i, 2, :] = vertices['c']
        array_subareas[i, 3, :] = vertices['d']

        # TODO: array_panels can be simplified (now only vector)
        # array_panels[i, 0] = 1
        array_panels[i, 1] = entry.num_chordwise_panels
        i += 1

    # Override symmetry flags (0, do not mirror anything that is passed in)
    array_symmetry = np.zeros((num_p), dtype=int, order='C')

    lattice.p = np.zeros((num_p, 4, 3), dtype=float, order='C')
    lattice.v = np.zeros((num_p, 4, 3), dtype=float, order='C')
    lattice.c = np.zeros((num_p, 3), dtype=float, order='C')
    lattice.bound_leg_midpoints = np.zeros((num_p, 3), dtype=float, order='C')
    lattice.n = np.zeros((num_p, 3), dtype=float, order='C')
    lattice.a = np.zeros((num_p), dtype=float, order='C')
    lattice.epsilon = settings.settings['_epsilon']

    logger.info("Generating lattice...")
    c_vlm.py2c_lattice(lattice, state, array_subareas, array_symmetry, array_panels)

    # ----- Compute the length of the bound leg -----
    logger.info(f"--> Number of panels: {lattice.info['num_panels']}")
    logger.info(f"--> Min panel area = {lattice.info['area_min']:.3e}")
    logger.info(f"--> Max panel area = {lattice.info['area_max']:.3e}")
    logger.info(f"--> Avg panel area = {lattice.info['area_avg']:.3e}")
    logger.info(f"--> Min panel aspect ratio = {lattice.info['aspect_min']:.3e}")
    logger.info(f"--> Max panel aspect ratio = {lattice.info['aspect_max']:.3e}")
    logger.info(f"--> Avg panel aspect ratio = {lattice.info['aspect_avg']:.3e}")

    # ========== ROTATE NORMALS ==========
    if settings.settings['_do_normal_rotations']:
        for entry in lattice.panel_bookkeeping:
            subarea = entry.subarea
            pan_idx = entry.pan_idx
            mirror = entry.mirror

            # CONTROL SURFACE DEFLECTIONS
            if subarea.parent_control is not None:
                hinge_axis = subarea.abs_hinge_axis(mirror)

                if mirror:
                    deflection = subarea.parent_control.deflection_mirror
                else:
                    deflection = subarea.parent_control.deflection

                # If deflection is 0, do not attempt rotation
                if deflection:
                    deflection = np.deg2rad(deflection)
                    R = axis_rot_matrix(hinge_axis, deflection)

                    for i in pan_idx:
                        lattice.n[i, :] = R @ lattice.n[i, :]

            # CAMBER LINE
            eta_a = subarea.parent_subdivision.rel_vertices['eta_a']
            eta_b = subarea.parent_subdivision.rel_vertices['eta_b']
            eta_m = (eta_a + eta_b)/2
            airfoil = subarea.parent_segment.segment_airfoil.at_eta(eta_m)

            num_panels = len([i for i in pan_idx])
            collocation_xsi = subarea.get_xsi_for_collocation_points(num_panels)

            for pan_of_subarea, i in enumerate(pan_idx):
                rot_axis = subarea.abs_camber_line_rot_axis(mirror)
                xsi = collocation_xsi[pan_of_subarea]
                angle = np.deg2rad(airfoil.camber_line_angle(xsi))

                # If deflection is 0, do not attempt rotation
                if angle:
                    R = axis_rot_matrix(rot_axis, angle)
                    lattice.n[i, :] = R @ lattice.n[i, :]
    return lattice

def calc_downwash(lattice, vlmdata):
    """
    Generate downwash factors for aircraft lattice.

    Pre-allocate memory for the (num_p x num_p) downwash factor matrix.
    The downwash calculations are performed in C, directly into this matrix.

    Display matrix condition number in console and log file.

    Args:
        :lattice: (object) data structure for VLM lattice
        :vlmdata: (object) data structure for VLM input and output
    """

    logger.info("Pre-allocating downwash matrix in memory...")
    num_p = lattice.info['num_panels']
    vlmdata.matrix_downwash = np.zeros((num_p, num_p), dtype=float, order='C')

    logger.info("Computing downwash factors...")
    c_vlm.py2c_downwash(lattice, vlmdata.matrix_downwash)
    logger.info(f"--> Condition number = {np.linalg.cond(vlmdata.matrix_downwash):.3e}")

def calc_boundary(lattice, state, vlmdata):
    """
    Generate boundary conditions (RHS term) for VLM.

    Pre-allocate memory for the (num_p x 1) right-hand-side array.
    The right-hand side terms are computed in C, directly into this memory.

    Args:
        :lattice: (object) data structure for VLM lattice
        :state: (object) data structure for flight state
        :vlmdata: (object) data structure for VLM input and output
    """

    logger.info("Pre-allocating rhs array in memory...")
    num_p = lattice.info['num_panels']
    vlmdata.array_rhs = np.zeros((num_p), dtype=float, order='C')

    logger.info("Computing right-hand side term...")
    c_vlm.py2c_boundary(lattice, state, vlmdata.array_rhs)
    vlmdata.array_rhs = np.array(vlmdata.array_rhs)

def solver(vlmdata):
    """
    Solve linear system for vortex strengths

    Args:
        :vlmdata: (object) data structure for VLM input and output
    """

    logger.info("Solving linear system...")
    vlmdata.matrix_lu, vlmdata.array_pivots, vlmdata.panelwise['gamma'], _ \
        = lapack.dgesv(vlmdata.matrix_downwash, vlmdata.array_rhs)

########
########
########
    # Needed???
    # vlmdata.panelwise['gamma'] = np.array(vlmdata.panelwise['gamma'], dtype=float, order='C')
########
########
########

def calc_results(lattice, state, vlmdata):
    """
    Calculate inwash at collocation points

    Args:
        :lattice: (object) data structure for VLM lattice
        :state: (object) data structure for flight state
        :vlmdata: (object) data structure for VLM input and output
    """

    logger.info("Pre-allocating vortex-lattice method results...")
    num_p = lattice.info['num_panels']

    # Allocate memory for panelwise results
    for key in ['vx', 'vy', 'vz', 'vmag',  'fx', 'fy', 'fz', 'fmag', 'cp']:
        vlmdata.panelwise[key] = np.zeros((num_p), dtype=float, order='C')

    logger.info("Computing results...")
    c_vlm.py2c_results(lattice, state, vlmdata)

    logger.info(f"--> Fx = {vlmdata.forces['x']:10.3e}")
    logger.info(f"--> Fy = {vlmdata.forces['y']:10.3e}")
    logger.info(f"--> Fz = {vlmdata.forces['z']:10.3e}")
    logger.info(f"--> FD = {vlmdata.forces['D']:10.3e}")
    logger.info(f"--> FC = {vlmdata.forces['C']:10.3e}")
    logger.info(f"--> FL = {vlmdata.forces['L']:10.3e}")
    logger.info(f"--> Mx = {vlmdata.forces['l']:10.3e}")
    logger.info(f"--> My = {vlmdata.forces['m']:10.3e}")
    logger.info(f"--> Mz = {vlmdata.forces['n']:10.3e}")

    logger.info(f"--> Cx = {vlmdata.coeffs['x']:7.4f}")
    logger.info(f"--> Cy = {vlmdata.coeffs['y']:7.4f}")
    logger.info(f"--> Cz = {vlmdata.coeffs['z']:7.4f}")
    logger.info(f"--> CD = {vlmdata.coeffs['D']:7.4f}")
    logger.info(f"--> CC = {vlmdata.coeffs['C']:7.4f}")
    logger.info(f"--> CL = {vlmdata.coeffs['L']:7.4f}")
    logger.info(f"--> Cl = {vlmdata.coeffs['l']:7.4f}")
    logger.info(f"--> Cm = {vlmdata.coeffs['m']:7.4f}")
    logger.info(f"--> Cn = {vlmdata.coeffs['n']:7.4f}")