import sys
import copy
import math
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm

sys.path.append('../..')
import modules.commons as commons
import modules.modules as modules
import modules.attentions as attentions

from modules.commons import init_weights, get_padding

from modules.ddsp import mlp, gru, scale_function, remove_above_nyquist, upsample
from modules.ddsp import harmonic_synth, amp_to_impulse_response, fft_convolve
from modules.ddsp import resample
import utils

from modules.stft import TorchSTFT

import torch.distributions as D

from modules.losses import (
    generator_loss,
    discriminator_loss,
    feature_loss,
    kl_loss
)

LRELU_SLOPE = 0.1

class PostF0Decoder(nn.Module):
    def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, spk_channels=0):
        super().__init__()

        self.in_channels = in_channels
        self.filter_channels = filter_channels
        self.kernel_size = kernel_size
        self.p_dropout = p_dropout
        self.gin_channels = spk_channels

        self.drop = nn.Dropout(p_dropout)
        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)
        self.norm_1 = modules.LayerNorm(filter_channels)
        self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)
        self.norm_2 = modules.LayerNorm(filter_channels)
        self.proj = nn.Conv1d(filter_channels, 1, 1)

        if spk_channels != 0:
            self.cond = nn.Conv1d(spk_channels, in_channels, 1)

    def forward(self, x, x_mask, g=None):
        x = torch.detach(x)
        if g is not None:
            g = torch.detach(g)
            x = x + self.cond(g)
        x = self.conv_1(x * x_mask)
        x = torch.relu(x)
        x = self.norm_1(x)
        x = self.drop(x)
        x = self.conv_2(x * x_mask)
        x = torch.relu(x)
        x = self.norm_2(x)
        x = self.drop(x)
        x = self.proj(x * x_mask)
        return x * x_mask

class TextEncoder(nn.Module):
    def __init__(self,
                 c_dim,
                 out_channels,
                 hidden_channels,
                 filter_channels,
                 n_heads,
                 n_layers,
                 kernel_size,
                 p_dropout):
        super().__init__()
        self.out_channels = out_channels
        self.hidden_channels = hidden_channels
        self.filter_channels = filter_channels
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.kernel_size = kernel_size
        self.p_dropout = p_dropout

        self.pre_net = torch.nn.Linear(c_dim, hidden_channels)

        self.encoder = attentions.Encoder(
            hidden_channels,
            filter_channels,
            n_heads,
            n_layers,
            kernel_size,
            p_dropout)
        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)

    def forward(self, x, x_lengths):
        x = x.transpose(1,-1)
        x = self.pre_net(x)
        x = torch.transpose(x, 1, -1)  # [b, h, t]
        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
        x = self.encoder(x * x_mask, x_mask)
        x = self.proj(x) * x_mask
        return x, x_mask

def pad_v2(input_ele, mel_max_length=None):
    if mel_max_length:
        max_len = mel_max_length
    else:
        max_len = max([input_ele[i].size(0) for i in range(len(input_ele))])

    out_list = list()
    for i, batch in enumerate(input_ele):
        if len(batch.shape) == 1:
            one_batch_padded = F.pad(
                batch, (0, max_len - batch.size(0)), "constant", 0.0
            )
        elif len(batch.shape) == 2:
            one_batch_padded = F.pad(
                batch, (0, 0, 0, max_len - batch.size(0)), "constant", 0.0
            )
        out_list.append(one_batch_padded)
    out_padded = torch.stack(out_list)
    return out_padded

class LengthRegulator(nn.Module):
    """ Length Regulator """

    def __init__(self):
        super(LengthRegulator, self).__init__()

    def LR(self, x, duration, max_len):
        x = torch.transpose(x, 1, 2)
        output = list()
        mel_len = list()
        for batch, expand_target in zip(x, duration):
            expanded = self.expand(batch, expand_target)
            output.append(expanded)
            mel_len.append(expanded.shape[0])

        if max_len is not None:
            output = pad_v2(output, max_len)
        else:
            output = pad_v2(output)
        output = torch.transpose(output, 1, 2)
        return output, torch.LongTensor(mel_len)

    def expand(self, batch, predicted):
        predicted = torch.squeeze(predicted)
        out = list()

        for i, vec in enumerate(batch):
            expand_size = predicted[i].item()
            state_info_index = torch.unsqueeze(torch.arange(0, expand_size), 1).float()
            state_info_length = torch.unsqueeze(torch.Tensor([expand_size] * expand_size), 1).float()
            state_info = torch.cat([state_info_index, state_info_length], 1).to(vec.device)
            new_vec = vec.expand(max(int(expand_size), 0), -1)
            new_vec = torch.cat([new_vec, state_info], 1)
            out.append(new_vec)
        out = torch.cat(out, 0)
        return out

    def forward(self, x, duration, max_len):
        output, mel_len = self.LR(x, duration, max_len)
        return output, mel_len

class PriorDecoder(nn.Module):
    def __init__(self,
                 out_bn_channels,
                 hidden_channels,
                 filter_channels,
                 n_heads,
                 n_layers,
                 kernel_size,
                 p_dropout,
                 n_speakers=0,
                 spk_channels=0):
        super().__init__()
        self.out_bn_channels = out_bn_channels
        self.hidden_channels = hidden_channels
        self.filter_channels = filter_channels
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.kernel_size = kernel_size
        self.p_dropout = p_dropout
        self.spk_channels = spk_channels

        self.prenet = nn.Conv1d(hidden_channels , hidden_channels, 3, padding=1)
        self.decoder = attentions.FFT(
            hidden_channels,
            filter_channels,
            n_heads,
            n_layers,
            kernel_size,
            p_dropout)
        self.proj = nn.Conv1d(hidden_channels, out_bn_channels, 1)

        if n_speakers != 0:
            self.cond = nn.Conv1d(spk_channels, hidden_channels, 1)

    def forward(self, x, x_lengths, spk_emb=None):
        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)

        x = self.prenet(x) * x_mask

        if (spk_emb is not None):
            x = x + self.cond(spk_emb)

        x = self.decoder(x * x_mask, x_mask)

        bn = self.proj(x) * x_mask

        return bn, x_mask

class Decoder(nn.Module):
    def __init__(self,
                 out_channels,
                 hidden_channels,
                 filter_channels,
                 n_heads,
                 n_layers,
                 kernel_size,
                 p_dropout,
                 n_speakers=0,
                 spk_channels=0,
                 in_channels=None):
        super().__init__()
        self.out_channels = out_channels
        self.hidden_channels = hidden_channels
        self.filter_channels = filter_channels
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.kernel_size = kernel_size
        self.p_dropout = p_dropout
        self.spk_channels = spk_channels

        self.prenet = nn.Conv1d(in_channels if in_channels is not None else hidden_channels, hidden_channels, 3, padding=1)
        self.decoder = attentions.FFT(
            hidden_channels,
            filter_channels,
            n_heads,
            n_layers,
            kernel_size,
            p_dropout)
        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)

        if n_speakers != 0:
            self.cond = nn.Conv1d(spk_channels, hidden_channels, 1)

    def forward(self, x, x_lengths, spk_emb=None):
        x = torch.detach(x)
        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)

        x = self.prenet(x) * x_mask

        if (spk_emb is not None):
            x = x + self.cond(spk_emb)

        x = self.decoder(x * x_mask, x_mask)

        x = self.proj(x) * x_mask

        return x, x_mask

class F0Decoder(nn.Module):
    def __init__(self,
                 out_channels,
                 hidden_channels,
                 filter_channels,
                 n_heads,
                 n_layers,
                 kernel_size,
                 p_dropout,
                 n_speakers=0,
                 spk_channels=0,
                 in_channels=None):
        super().__init__()
        self.out_channels = out_channels
        self.hidden_channels = hidden_channels
        self.filter_channels = filter_channels
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.kernel_size = kernel_size
        self.p_dropout = p_dropout
        self.spk_channels = spk_channels

        self.prenet = nn.Conv1d(in_channels if in_channels is not None else hidden_channels, hidden_channels, 3, padding=1)
        self.decoder = attentions.FFT(
            hidden_channels,
            filter_channels,
            n_heads,
            n_layers,
            kernel_size,
            p_dropout)
        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
        self.f0_prenet = nn.Conv1d(1, hidden_channels , 3, padding=1)

        if n_speakers != 0:
            self.cond = nn.Conv1d(spk_channels, hidden_channels, 1)

    def forward(self, x, norm_f0, x_lengths, spk_emb=None):
        x = torch.detach(x)
        x += self.f0_prenet(norm_f0)
        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)

        x = self.prenet(x) * x_mask

        if (spk_emb is not None):
            x = x + self.cond(spk_emb)

        x = self.decoder(x * x_mask, x_mask)

        x = self.proj(x) * x_mask

        return x, x_mask

class ConvReluNorm(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.n_layers = n_layers
        self.p_dropout = p_dropout
        assert n_layers > 1, "Number of layers should be larger than 0."

        self.conv_layers = nn.ModuleList()
        self.norm_layers = nn.ModuleList()
        self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
        self.norm_layers.append(LayerNorm(hidden_channels))
        self.relu_drop = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(p_dropout))
        for _ in range(n_layers - 1):
            self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
            self.norm_layers.append(LayerNorm(hidden_channels))
        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
        self.proj.weight.data.zero_()
        self.proj.bias.data.zero_()

    def forward(self, x):
        x = self.conv_layers[0](x)
        x = self.norm_layers[0](x)
        x = self.relu_drop(x)

        for i in range(1, self.n_layers):
            x_ = self.conv_layers[i](x)
            x_ = self.norm_layers[i](x_)
            x_ = self.relu_drop(x_)
            x = (x + x_) / 2
        x = self.proj(x)
        return x

class PosteriorEncoder(nn.Module):
    def __init__(self,
                 hps,
                 in_channels,
                 out_channels,
                 hidden_channels,
                 kernel_size,
                 dilation_rate,
                 n_layers):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers

        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
        self.enc = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, n_speakers=hps.data.n_speakers, spk_channels=hps.model.spk_channels)
        # self.enc = ConvReluNorm(hidden_channels,
        #                         hidden_channels,
        #                         hidden_channels,
        #                         kernel_size,
        #                         n_layers,
        #                         0.1)
        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)

    def forward(self, x, x_lengths, g=None):
        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
        x = self.pre(x) * x_mask
        x = self.enc(x, x_mask, g=g)
        stats = self.proj(x) * x_mask
        return stats, x_mask

class ResBlock3(torch.nn.Module):
    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
        super(ResBlock3, self).__init__()
        self.convs = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0])))
        ])
        self.convs.apply(init_weights)

    def forward(self, x, x_mask=None):
        for c in self.convs:
            xt = F.leaky_relu(x, LRELU_SLOPE)
            if x_mask is not None:
                xt = xt * x_mask
            xt = c(xt)
            x = xt + x
        if x_mask is not None:
            x = x * x_mask
        return x

    def remove_weight_norm(self):
        for l in self.convs:
            remove_weight_norm(l)

class Generator_Harm(torch.nn.Module):
    def __init__(self, hps):
        super(Generator_Harm, self).__init__()
        self.hps = hps

        self.prenet = Conv1d(hps.model.hidden_channels, hps.model.hidden_channels, 3, padding=1)

        self.net = ConvReluNorm(hps.model.hidden_channels,
                                hps.model.hidden_channels,
                                hps.model.hidden_channels,
                                hps.model.kernel_size,
                                8,
                                hps.model.p_dropout)

        # self.rnn = nn.LSTM(input_size=hps.model.hidden_channels,
        #    hidden_size=hps.model.hidden_channels,
        #    num_layers=1,
        #    bias=True,
        #    batch_first=True,
        #    dropout=0.5,
        #    bidirectional=True)
        self.postnet = Conv1d(hps.model.hidden_channels, hps.model.n_harmonic + 1, 3, padding=1)

    def forward(self, f0, harm, mask):
        pitch = f0.transpose(1, 2)
        harm = self.prenet(harm)

        harm = self.net(harm) * mask
        # harm = harm.transpose(1, 2)
        # harm, (hs, hc) = self.rnn(harm)
        # harm = harm.transpose(1, 2)

        harm = self.postnet(harm)
        harm = harm.transpose(1, 2)
        param = harm

        param = scale_function(param)
        total_amp = param[..., :1]
        amplitudes = param[..., 1:]
        amplitudes = remove_above_nyquist(
            amplitudes,
            pitch,
            self.hps.data.sampling_rate,
        )
        amplitudes /= amplitudes.sum(-1, keepdim=True)
        amplitudes *= total_amp

        amplitudes = upsample(amplitudes, self.hps.data.hop_length)
        pitch = upsample(pitch, self.hps.data.hop_length)

        n_harmonic = amplitudes.shape[-1]
        omega = torch.cumsum(2 * math.pi * pitch / self.hps.data.sampling_rate, 1)
        omegas = omega * torch.arange(1, n_harmonic + 1).to(omega)
        signal_harmonics = (torch.sin(omegas) * amplitudes)
        signal_harmonics = signal_harmonics.transpose(1, 2)
        return signal_harmonics

class Generator(torch.nn.Module):
    def __init__(self, hps, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates,
                 upsample_initial_channel, upsample_kernel_sizes, n_speakers=0, spk_channels=0):
        super(Generator, self).__init__()
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
        self.upsample_rates = upsample_rates
        self.n_speakers = n_speakers

        resblock = modules.ResBlock1 if resblock == '1' else modules.R

        self.downs = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            i = len(upsample_rates) - 1 - i
            u = upsample_rates[i]
            k = upsample_kernel_sizes[i]
            # print("down: ",upsample_initial_channel//(2**(i+1))," -> ", upsample_initial_channel//(2**i))
            self.downs.append(weight_norm(
                Conv1d(hps.model.n_harmonic + 2, hps.model.n_harmonic + 2,
                       k, u, padding=k // 2)))

        self.resblocks_downs = nn.ModuleList()
        for i in range(len(self.downs)):
            j = len(upsample_rates) - 1 - i
            self.resblocks_downs.append(ResBlock3(hps.model.n_harmonic + 2, 3, (1, 3)))

        self.concat_pre = Conv1d(upsample_initial_channel + hps.model.n_harmonic + 2, upsample_initial_channel, 3, 1,
                                 padding=1)
        self.concat_conv = nn.ModuleList()
        for i in range(len(upsample_rates)):
            ch = upsample_initial_channel // (2 ** (i + 1))
            self.concat_conv.append(Conv1d(ch + hps.model.n_harmonic + 2, ch, 3, 1, padding=1, bias=False))

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(weight_norm(
                ConvTranspose1d(upsample_initial_channel // (2 ** i), upsample_initial_channel // (2 ** (i + 1)),
                                k, u, padding=(k - u) // 2)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel // (2 ** (i + 1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
        self.ups.apply(init_weights)

        if self.n_speakers != 0:
            self.cond = nn.Conv1d(spk_channels, upsample_initial_channel, 1)

    def forward(self, x, ddsp, g=None):

        x = self.conv_pre(x)

        if g is not None:
            x = x + self.cond(g)

        se = ddsp
        res_features = [se]
        for i in range(self.num_upsamples):
            in_size = se.size(2)
            se = self.downs[i](se)
            se = self.resblocks_downs[i](se)
            up_rate = self.upsample_rates[self.num_upsamples - 1 - i]
            se = se[:, :, : in_size // up_rate]
            res_features.append(se)

        x = torch.cat([x, se], 1)
        x = self.concat_pre(x)

        for i in range(self.num_upsamples):
            x = F.leaky_relu(x, modules.LRELU_SLOPE)
            in_size = x.size(2)
            x = self.ups[i](x)
            # 保证维度正确，丢掉多余通道
            x = x[:, :, : in_size * self.upsample_rates[i]]

            x = torch.cat([x, res_features[self.num_upsamples - 1 - i]], 1)
            x = self.concat_conv[i](x)

            xs = None
            for j in range(self.num_kernels):
                if xs is None:
                    xs = self.resblocks[i * self.num_kernels + j](x)
                else:
                    xs += self.resblocks[i * self.num_kernels + j](x)
            x = xs / self.num_kernels

        x = F.leaky_relu(x)
        x = self.conv_post(x)
        x = torch.tanh(x)

        return x

    def remove_weight_norm(self):
        print('Removing weight norm...')
        for l in self.ups:
            remove_weight_norm(l)
        for l in self.resblocks:
            l.remove_weight_norm()

class Generator_Noise(torch.nn.Module):
    def __init__(self, hps):
        super(Generator_Noise, self).__init__()
        self.hps = hps
        self.win_size = hps.data.win_size
        self.hop_size = hps.data.hop_length
        self.fft_size = hps.data.n_fft
        self.istft_pre = Conv1d(hps.model.hidden_channels, hps.model.hidden_channels, 3, padding=1)

        self.net = ConvReluNorm(hps.model.hidden_channels,
                                hps.model.hidden_channels,
                                hps.model.hidden_channels,
                                hps.model.kernel_size,
                                8,
                                hps.model.p_dropout)

        self.istft_amplitude = torch.nn.Conv1d(hps.model.hidden_channels, self.fft_size // 2 + 1, 1, 1)
        self.window = torch.hann_window(self.win_size)

    def forward(self, x, mask):
        istft_x = x
        istft_x = self.istft_pre(istft_x)

        istft_x = self.net(istft_x) * mask

        amp = self.istft_amplitude(istft_x).unsqueeze(-1)
        phase = (torch.rand(amp.shape) * 2 * 3.14 - 3.14).to(amp)

        real = amp * torch.cos(phase)
        imag = amp * torch.sin(phase)
        spec = torch.cat([real, imag], 3)
        istft_x = torch.istft(spec, self.fft_size, self.hop_size, self.win_size, self.window.to(amp), True,
                              length=x.shape[2] * self.hop_size, return_complex=False)

        return istft_x.unsqueeze(1)

class LayerNorm(nn.Module):
    def __init__(self, channels, eps=1e-5):
        super().__init__()
        self.channels = channels
        self.eps = eps

        self.gamma = nn.Parameter(torch.ones(channels))
        self.beta = nn.Parameter(torch.zeros(channels))

    def forward(self, x):
        x = x.transpose(1, -1)
        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
        return x.transpose(1, -1)

class DiscriminatorP(torch.nn.Module):
    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
        super(DiscriminatorP, self).__init__()
        self.period = period
        self.use_spectral_norm = use_spectral_norm
        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
        self.convs = nn.ModuleList([
            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))),
        ])
        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))

    def forward(self, x):
        fmap = []

        # 1d to 2d
        b, c, t = x.shape
        if t % self.period != 0:  # pad first
            n_pad = self.period - (t % self.period)
            x = F.pad(x, (0, n_pad), "reflect")
            t = t + n_pad
        x = x.view(b, c, t // self.period, self.period)

        for l in self.convs:
            x = l(x)
            x = F.leaky_relu(x, modules.LRELU_SLOPE)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        x = torch.flatten(x, 1, -1)

        return x, fmap

class DiscriminatorS(torch.nn.Module):
    def __init__(self, use_spectral_norm=False):
        super(DiscriminatorS, self).__init__()
        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
        self.convs = nn.ModuleList([
            norm_f(Conv1d(1, 16, 15, 1, padding=7)),
            norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
            norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
            norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
            norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
        ])
        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))

    def forward(self, x):
        fmap = []

        for l in self.convs:
            x = l(x)
            x = F.leaky_relu(x, modules.LRELU_SLOPE)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        x = torch.flatten(x, 1, -1)

        return x, fmap

class MultiFrequencyDiscriminator(nn.Module):
    def __init__(self,
                 hop_lengths=[128, 256, 512],
                 hidden_channels=[256, 512, 512],
                 domain='double', mel_scale=True):
        super(MultiFrequencyDiscriminator, self).__init__()

        self.stfts = nn.ModuleList([
            TorchSTFT(fft_size=x * 4, hop_size=x, win_size=x * 4,
                      normalized=True, domain=domain, mel_scale=mel_scale)
            for x in hop_lengths])

        self.domain = domain
        if domain == 'double':
            self.discriminators = nn.ModuleList([
                BaseFrequenceDiscriminator(2, c)
                for x, c in zip(hop_lengths, hidden_channels)])
        else:
            self.discriminators = nn.ModuleList([
                BaseFrequenceDiscriminator(1, c)
                for x, c in zip(hop_lengths, hidden_channels)])

    def forward(self, x):
        scores, feats = list(), list()
        for stft, layer in zip(self.stfts, self.discriminators):
            # print(stft)
            mag, phase = stft.transform(x.squeeze())
            if self.domain == 'double':
                mag = torch.stack(torch.chunk(mag, 2, dim=1), dim=1)
            else:
                mag = mag.unsqueeze(1)

            score, feat = layer(mag)
            scores.append(score)
            feats.append(feat)
        return scores, feats

class BaseFrequenceDiscriminator(nn.Module):
    def __init__(self, in_channels, hidden_channels=512):
        super(BaseFrequenceDiscriminator, self).__init__()

        self.discriminator = nn.ModuleList()
        self.discriminator += [
            nn.Sequential(
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    in_channels, hidden_channels // 32,
                    kernel_size=(3, 3), stride=(1, 1)))
            ),
            nn.Sequential(
                nn.LeakyReLU(0.2, True),
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    hidden_channels // 32, hidden_channels // 16,
                    kernel_size=(3, 3), stride=(2, 2)))
            ),
            nn.Sequential(
                nn.LeakyReLU(0.2, True),
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    hidden_channels // 16, hidden_channels // 8,
                    kernel_size=(3, 3), stride=(1, 1)))
            ),
            nn.Sequential(
                nn.LeakyReLU(0.2, True),
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    hidden_channels // 8, hidden_channels // 4,
                    kernel_size=(3, 3), stride=(2, 2)))
            ),
            nn.Sequential(
                nn.LeakyReLU(0.2, True),
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    hidden_channels // 4, hidden_channels // 2,
                    kernel_size=(3, 3), stride=(1, 1)))
            ),
            nn.Sequential(
                nn.LeakyReLU(0.2, True),
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    hidden_channels // 2, hidden_channels,
                    kernel_size=(3, 3), stride=(2, 2)))
            ),
            nn.Sequential(
                nn.LeakyReLU(0.2, True),
                nn.ReflectionPad2d((1, 1, 1, 1)),
                nn.utils.weight_norm(nn.Conv2d(
                    hidden_channels, 1,
                    kernel_size=(3, 3), stride=(1, 1)))
            )
        ]

    def forward(self, x):
        hiddens = []
        for layer in self.discriminator:
            x = layer(x)
            hiddens.append(x)
        return x, hiddens[-1]

class Discriminator(torch.nn.Module):
    def __init__(self, hps, use_spectral_norm=False):
        super(Discriminator, self).__init__()
        periods = [2, 3, 5, 7, 11]

        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
        discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
        self.discriminators = nn.ModuleList(discs)
        # self.disc_multfrequency = MultiFrequencyDiscriminator(hop_lengths=[int(hps.data.sampling_rate * 2.5 / 1000),
        #                                                                    int(hps.data.sampling_rate * 5 / 1000),
        #                                                                    int(hps.data.sampling_rate * 7.5 / 1000),
        #                                                                    int(hps.data.sampling_rate * 10 / 1000),
        #                                                                    int(hps.data.sampling_rate * 12.5 / 1000),
        #                                                                    int(hps.data.sampling_rate * 15 / 1000)],
        #                                                       hidden_channels=[256, 256, 256, 256, 256])

    def forward(self, y, y_hat):
        y_d_rs = []
        y_d_gs = []
        fmap_rs = []
        fmap_gs = []
        for i, d in enumerate(self.discriminators):
            y_d_r, fmap_r = d(y)
            y_d_g, fmap_g = d(y_hat)
            y_d_rs.append(y_d_r)
            y_d_gs.append(y_d_g)
            fmap_rs.append(fmap_r)
            fmap_gs.append(fmap_g)
        # scores_r, fmaps_r = self.disc_multfrequency(y)
        # scores_g, fmaps_g = self.disc_multfrequency(y_hat)
        # for i in range(len(scores_r)):
        #     y_d_rs.append(scores_r[i])
        #     y_d_gs.append(scores_g[i])
        #     fmap_rs.append(fmaps_r[i])
        #     fmap_gs.append(fmaps_g[i])
        return y_d_rs, y_d_gs, fmap_rs, fmap_gs

class SynthesizerTrn(nn.Module):
    """
    Model
    """

    def __init__(self, hps):
        super().__init__()
        self.hps = hps

        self.text_encoder = TextEncoder(
            hps.data.c_dim,
            hps.model.prior_hidden_channels,
            hps.model.prior_hidden_channels,
            hps.model.prior_filter_channels,
            hps.model.prior_n_heads,
            hps.model.prior_n_layers,
            hps.model.prior_kernel_size,
            hps.model.prior_p_dropout)

        self.decoder = PriorDecoder(
            hps.model.hidden_channels * 2,
            hps.model.prior_hidden_channels,
            hps.model.prior_filter_channels,
            hps.model.prior_n_heads,
            hps.model.prior_n_layers,
            hps.model.prior_kernel_size,
            hps.model.prior_p_dropout,
            n_speakers=hps.data.n_speakers,
            spk_channels=hps.model.spk_channels
        )

        self.f0_decoder = F0Decoder(
            1,
            hps.model.prior_hidden_channels,
            hps.model.prior_filter_channels,
            hps.model.prior_n_heads,
            hps.model.prior_n_layers,
            hps.model.prior_kernel_size,
            hps.model.prior_p_dropout,
            n_speakers=hps.data.n_speakers,
            spk_channels=hps.model.spk_channels
        )

        self.mel_decoder = Decoder(
            hps.data.acoustic_dim,
            hps.model.prior_hidden_channels,
            hps.model.prior_filter_channels,
            hps.model.prior_n_heads,
            hps.model.prior_n_layers,
            hps.model.prior_kernel_size,
            hps.model.prior_p_dropout,
            n_speakers=hps.data.n_speakers,
            spk_channels=hps.model.spk_channels
        )

        self.posterior_encoder = PosteriorEncoder(
            hps,
            hps.data.acoustic_dim,
            hps.model.hidden_channels,
            hps.model.hidden_channels, 3, 1, 8)

        self.dropout = nn.Dropout(0.2)

        self.LR = LengthRegulator()

        self.dec = Generator(hps,
                             hps.model.hidden_channels,
                             hps.model.resblock,
                             hps.model.resblock_kernel_sizes,
                             hps.model.resblock_dilation_sizes,
                             hps.model.upsample_rates,
                             hps.model.upsample_initial_channel,
                             hps.model.upsample_kernel_sizes,
                             n_speakers=hps.data.n_speakers,
                             spk_channels=hps.model.spk_channels)

        self.dec_harm = Generator_Harm(hps)

        self.dec_noise = Generator_Noise(hps)

        self.f0_prenet = nn.Conv1d(1, hps.model.prior_hidden_channels , 3, padding=1)
        self.energy_prenet = nn.Conv1d(1, hps.model.prior_hidden_channels , 3, padding=1)
        self.mel_prenet = nn.Conv1d(hps.data.acoustic_dim, hps.model.prior_hidden_channels , 3, padding=1)

        if hps.data.n_speakers > 0:
            self.emb_spk = nn.Embedding(hps.data.n_speakers, hps.model.spk_channels)
        self.flow = modules.ResidualCouplingBlock(hps.model.prior_hidden_channels, hps.model.hidden_channels, 5, 1, 4,n_speakers=hps.data.n_speakers, gin_channels=hps.model.spk_channels)

    def forward(self, c, c_lengths, F0, uv, mel, bn_lengths, spk_id=None):
        if self.hps.data.n_speakers > 0:
            g = self.emb_spk(spk_id).unsqueeze(-1)  # [b, h, 1]
        else:
            g = None

        # Encoder
        decoder_input, x_mask = self.text_encoder(c, c_lengths)

        LF0 = 2595. * torch.log10(1. + F0 / 700.)
        LF0 = LF0 / 500
        norm_f0 = utils.normalize_f0(LF0,x_mask, uv.squeeze(1),random_scale=True)
        pred_lf0, predict_bn_mask = self.f0_decoder(decoder_input, norm_f0, bn_lengths, spk_emb=g)
        # print(pred_lf0)
        loss_f0 = F.mse_loss(pred_lf0, LF0)

        # aam
        predict_mel, predict_bn_mask = self.mel_decoder(decoder_input + self.f0_prenet(LF0), bn_lengths, spk_emb=g)

        predict_energy = predict_mel.detach().sum(1).unsqueeze(1) / self.hps.data.acoustic_dim

        decoder_input = decoder_input + \
                        self.f0_prenet(LF0) + \
                        self.energy_prenet(predict_energy) + \
                        self.mel_prenet(predict_mel.detach())
        decoder_output, predict_bn_mask = self.decoder(decoder_input, bn_lengths, spk_emb=g)

        prior_info = decoder_output
        m_p = prior_info[:, :self.hps.model.hidden_channels, :]
        logs_p = prior_info[:, self.hps.model.hidden_channels:, :]

        # posterior
        posterior, y_mask = self.posterior_encoder(mel, bn_lengths,g=g)

        m_q = posterior[:, :self.hps.model.hidden_channels, :]
        logs_q = posterior[:, self.hps.model.hidden_channels:, :]
        z = (m_q + torch.randn_like(m_q) * torch.exp(logs_q)) * y_mask
        z_p = self.flow(z, y_mask, g=g)

        # kl loss
        loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, y_mask)

        p_z = z
        p_z = self.dropout(p_z)

        pitch = upsample(F0.transpose(1, 2), self.hps.data.hop_length)
        omega = torch.cumsum(2 * math.pi * pitch / self.hps.data.sampling_rate, 1)
        sin = torch.sin(omega).transpose(1, 2)

        # dsp synthesize
        noise_x = self.dec_noise(p_z, y_mask)
        harm_x = self.dec_harm(F0, p_z, y_mask)

        # dsp waveform
        dsp_o = torch.cat([harm_x, noise_x], axis=1)

        decoder_condition = torch.cat([harm_x, noise_x, sin], axis=1)

        # dsp based HiFiGAN vocoder
        x_slice, ids_slice = commons.rand_slice_segments(p_z, bn_lengths,
                                                         self.hps.train.segment_size // self.hps.data.hop_length)
        F0_slice = commons.slice_segments(F0, ids_slice, self.hps.train.segment_size // self.hps.data.hop_length)
        dsp_slice = commons.slice_segments(dsp_o, ids_slice * self.hps.data.hop_length, self.hps.train.segment_size)
        condition_slice = commons.slice_segments(decoder_condition, ids_slice * self.hps.data.hop_length,
                                                 self.hps.train.segment_size)
        o = self.dec(x_slice, condition_slice.detach(), g=g)

        return o, ids_slice, LF0 * predict_bn_mask, dsp_slice.sum(1), loss_kl, \
               predict_mel, predict_bn_mask, pred_lf0, loss_f0, norm_f0

    def infer(self, c, g=None, f0=None,uv=None, predict_f0=False, noice_scale=0.3):
        if len(g.shape) == 2:
            g = g.squeeze(0)
        if len(f0.shape) == 2:
            f0 = f0.unsqueeze(0)
        g = self.emb_spk(g).unsqueeze(-1)  # [b, h, 1]

        c_lengths = (torch.ones(c.size(0)) * c.size(-1)).to(c.device)

        # Encoder
        decoder_input, x_mask = self.text_encoder(c, c_lengths)
        y_lengths = c_lengths

        LF0 = 2595. * torch.log10(1. + f0 / 700.)
        LF0 = LF0 / 500

        if predict_f0:
            norm_f0 = utils.normalize_f0(LF0, x_mask, uv.squeeze(1))
            pred_lf0, predict_bn_mask = self.f0_decoder(decoder_input, norm_f0, y_lengths, spk_emb=g)
            pred_f0 = 700 * ( torch.pow(10, pred_lf0 * 500 / 2595) - 1)
            f0 = pred_f0
            LF0 = pred_lf0

        # aam
        predict_mel, predict_bn_mask = self.mel_decoder(decoder_input + self.f0_prenet(LF0), y_lengths, spk_emb=g)
        predict_energy = predict_mel.sum(1).unsqueeze(1) / self.hps.data.acoustic_dim

        decoder_input = decoder_input + \
                        self.f0_prenet(LF0) + \
                        self.energy_prenet(predict_energy) + \
                        self.mel_prenet(predict_mel)
        decoder_output, y_mask = self.decoder(decoder_input, y_lengths, spk_emb=g)

        prior_info = decoder_output

        m_p = prior_info[:, :self.hps.model.hidden_channels, :]
        logs_p = prior_info[:, self.hps.model.hidden_channels:, :]
        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noice_scale
        z = self.flow(z_p, y_mask, g=g, reverse=True)

        prior_z = z

        noise_x = self.dec_noise(prior_z, y_mask)

        harm_x = self.dec_harm(f0, prior_z, y_mask)

        pitch = upsample(f0.transpose(1, 2), self.hps.data.hop_length)
        omega = torch.cumsum(2 * math.pi * pitch / self.hps.data.sampling_rate, 1)
        sin = torch.sin(omega).transpose(1, 2)

        decoder_condition = torch.cat([harm_x, noise_x, sin], axis=1)

        # dsp based HiFiGAN vocoder
        o = self.dec(prior_z, decoder_condition, g=g)

        return o, harm_x.sum(1).unsqueeze(1), noise_x, f0


import asyncio
from connectionBDD import *
from connectionZabbix import *
from Zapi import *
from datetime import datetime
import time
import sys
import pandas as pd
import os
from pandasql import sqldf
from pyzabbix import ZabbixAPI

#On récupère le path pour que le fichier exécutable soit créé dans le même répertoire
path = os.path.dirname(sys.executable)

#Pour mettre la date courante dans le nom du fichier csv crée avec un fstring
now = datetime.now()
nowstr = now.strftime("%d%m%Y")

"""
#Enabling debug logging for the pyzabbix API
stream = logging.StreamHandler(sys.stdout)
stream.setLevel(logging.DEBUG)
log = logging.getLogger('pyzabbix')
log.addHandler(stream)
log.setLevel(logging.DEBUG)
"""

async def main():

    sortie = 'preprod_sortie_ems'
    coordo = 'preprod_bdd_coordination'
    histo = 'preprod_historisation'

    #Connection aux 2 BDDs
    try:
        conn_sortie = ConnectionBDD(sortie)
        print(f"Connection to " +sortie+ " created successfully.")
    except Exception as ex:
        print("Connection to " +sortie+ " could not be made due to the following error: \n", ex)

    try:
        conn_coordo = ConnectionBDD(coordo, "bdd_coordination_schema")
        print(f"Connection to " +coordo+ " created successfully.")
    except Exception as ex:
        print("Connection to " +coordo+ " could not be made due to the following error: \n", ex)

    #Connection aux tables dont on a besoin
    try:
        conn_sortie.get_table('result')
        print(f"Table for result created successfully.")
    except Exception as ex:
        print("Table for result could not be made due to the following error: \n", ex)

    try:
        conn_sortie.get_table('equipement_pilote_consommation_moyenne')
        print(f"Table for conso created successfully.")
    except Exception as ex:
        print("Table for conso could not be made due to the following error: \n", ex)

    try:
        conn_sortie.get_table('p_c_with_flexible_consumption')
        print(f"Table for p_c_with_flexible_consumption created successfully.")
    except Exception as ex:
        print("Table for p_c_with_flexible_consumption could not be made due to the following error: \n", ex)

    try:
        conn_sortie.get_table('p_c_without_flexible_consumption')
        print(f"Table for p_c_without_flexible_consumption created successfully.")
    except Exception as ex:
        print("Table for p_c_without_flexible_consumption could not be made due to the following error: \n", ex)

    try:
        conn_coordo.get_table_with_schema('equipement_pilote_ou_mesure', 'bdd_coordination_schema')
        print(f"Table for equipement_pilote_ou_mesure created successfully.")
    except Exception as ex:
        print("Table for equipement_pilote_ou_mesure could not be made due to the following error: \n", ex)

    try:
        conn_coordo.get_table_with_schema('equipement_mesure_compteur_electrique', 'bdd_coordination_schema')
        print(f"Table for equipement_mesure_compteur_electrique created successfully.")
    except Exception as ex:
        print("Table for equipement_mesure_compteur_electrique could not be made due to the following error: \n", ex)

    #Calcul du nombre d'appareils connectés par 1/4h
    def indic_appareils_connectes() -> int :
        with conn_sortie.engine.connect() as conn:
            appconn = pd.read_sql("SELECT COUNT(*) FROM result\
                                    GROUP BY first_valid_timestamp\
                                    ORDER BY first_valid_timestamp DESC\
                                    LIMIT 1"\
                                    , con = conn)
        return int(appconn.loc[0]['count'])

    #Calcul du nombre d'appareils pilotés la semaine dernière
    def indic_appareils_pilotes_semaine() -> int :
        with conn_sortie.engine.connect() as conn:
            appsem = pd.read_sql("SELECT COUNT(DISTINCT machine_id) AS Nombre_appareils_pilotés_la_semaine_dernière FROM result\
                                    WHERE first_valid_timestamp > (CAST(EXTRACT (epoch FROM NOW()) AS INT) - 604800)", con = conn)
        return appsem

    #Calcul du pourcentage des appareils de ELFE lancés dans les dernières 24h
    # Actuellement, il est mis au niveau du pourcentage de machines activées dans Viriya
    #Rajouter des types dans la liste si de nouveaux types qu'on ne peut pas piloter s'ajoutent
    def pourcentage_app_lancés_24h() -> float :
        with conn_coordo.engine.connect() as conn:
            nb_app_EMS = pd.read_sql("SELECT COUNT(*) FROM equipement_pilote_ou_mesure\
                                        WHERE equipement_pilote_ou_mesure_type_id NOT IN (410, 901, 910, 920)"
                                    , con = conn)
        nb_app_EMS = int(nb_app_EMS.loc[0]['count'])
        with conn_sortie.engine.connect() as conn:
            nb_app_lancés_24h_continu = pd.read_sql("SELECT COUNT(DISTINCT machine_id) FROM result \
                                                        WHERE machine_type IN (131, -1)\
                                                        AND decisions_0 = 1\
                                                        AND first_valid_timestamp >= (CAST(EXTRACT (epoch FROM NOW()) AS INT) - 86400)"
                                        , con = conn) # Compléter les machine_type
            nb_app_lancés_24h_discontinu = pd.read_sql(""" SELECT COUNT(*) FROM result
                                                WHERE decisions_0 = 1
                                                AND machine_type IN (221)
                                                AND first_valid_timestamp > (CAST(EXTRACT (epoch FROM NOW()) AS INT) - 86400)"""
                                        , con = conn)
        nb_app_lancés_24h = int(nb_app_lancés_24h_continu.loc[0]['count']) + int(nb_app_lancés_24h_discontinu.loc[0]['count'])
        pourcentage_app_lancés_24h = (100*nb_app_lancés_24h)/nb_app_EMS
        return round(pourcentage_app_lancés_24h, 1)
        #Liste continu à terme (en prod) : 131, 151, 155
        #Liste discontinu à terme (en prod) : 515, 112, 111, 113, 221, 225

    #Calcul du cumul de kW d'énergie placés par l'EMS depuis le début du projet
    def cumul_enr() -> int:
        with conn_sortie.engine.connect() as conn:
            #On fabrique deux tables sous la forme machine_type | nombre de lancements
            #1e table pour les machines discoutinues : on compte chaque lancement
            coeffs_discontinu:pd.DataFrame = pd.read_sql("SELECT machine_type, COUNT(*) FROM result\
                                                WHERE decisions_0 = 1\
                                                AND machine_type IN (111)\
                                            GROUP BY machine_type"
                                    , con = conn) #Compléter les machines_type
            #2e table pour les machines continues : si pas de arrêt / relance, on compte un lancement par jour
            coeffs_continu:pd.DataFrame = pd.read_sql("""SELECT * FROM
                                                        (SELECT machine_type, COUNT(*) FROM
                                                            (SELECT DISTINCT * FROM
                                                                (SELECT first_valid_timestamp/86400 AS day, machine_id, machine_type FROM
                                                                        result WHERE decisions_0 = 1) AS T1) AS T2
                                                                        GROUP BY machine_type) AS T3 WHERE machine_type IN (131)"""
                                    , con = conn) #Compléter les machines_type
            #On récupère la table machine_type | consommation moyenne.
            with conn_coordo.engine.connect() as conn:
                conso_energie:pd.DataFrame = pd.read_sql("SELECT equipement_pilote_type_id, consommation FROM equipement_pilote_consommation_moyenne"
                        , con = conn)
            #cumul_enr correspond à l'indicateur final qu'on initialise à 0
            cumul_enr = 0
            #On l'incrémente avec pour chaque ligne des tableaux continus et discontinus nb de machines d'un type * moyenne de l'énergie consommée par ce type de machine
            for i in coeffs_discontinu.index:
                for j in conso_energie.index:
                    if coeffs_discontinu.loc[i]['machine_type'] == conso_energie.loc[j]['equipement_pilote_type_id']:
                        cumul_enr += coeffs_discontinu.loc[i]['count'] * conso_energie.loc[j]['consommation']
            for i in coeffs_continu.index:
                for j in conso_energie.index:
                    if coeffs_continu.loc[i]['machine_type'] == conso_energie.loc[j]['equipement_pilote_type_id']:
                        cumul_enr += coeffs_continu.loc[i]['count'] * conso_energie.loc[j]['consommation']
            print("Ancien calcul : ", int(cumul_enr/1000))
            return int(cumul_enr/1000)

    #Calcul du cumul de kW d'énergie placés par l'EMS depuis le début du projet
    def cumul_enr_opti() -> int:
        with conn_sortie.engine.connect() as conn:
            #On fabrique deux tables sous la forme machine_type | nombre de lancements
            #1e table pour les machines discoutinues : on compte chaque lancement
            coeffs_discontinu:pd.DataFrame = pd.read_sql("SELECT machine_type, COUNT(*) FROM result\
                                                WHERE decisions_0 = 1\
                                                AND machine_type IN (111)\
                                            GROUP BY machine_type"
                                    , con = conn) #Compléter les machines_type
            print("Coeffs discontinus : " , coeffs_discontinu)
            #2e table pour les machines continues : si pas de arrêt / relance, on compte un lancement par jour
            coeffs_continu:pd.DataFrame = pd.read_sql("""SELECT * FROM
                                                        (SELECT machine_type, COUNT(*) FROM
                                                            (SELECT DISTINCT * FROM
                                                                (SELECT first_valid_timestamp/86400 AS day, machine_id, machine_type FROM result
                                                                        WHERE decisions_0 = 1) AS T1) AS T2
                                                                        GROUP BY machine_type) AS T3 WHERE machine_type IN (131)"""
                                    , con = conn) #Compléter les machines_type
            print("Coeffs continus : ", coeffs_continu)

            #On récupère la table machine_type | consommation moyenne.
            with conn_coordo.engine.connect() as conn:
                conso_energie:pd.DataFrame = pd.read_sql("SELECT equipement_pilote_type_id, consommation FROM equipement_pilote_consommation_moyenne"
                        , con = conn)
            #cumul_enr correspond à l'indicateur final qu'on initialise à 0
            cumul_enr = 0
            #On l'incrémente avec pour chaque ligne des tableaux continus et discontinus nb de machines d'un type * moyenne de l'énergie consommée par ce type de machine
            for i in coeffs_discontinu.index:
                for j in conso_energie.index:
                    if coeffs_discontinu.loc[i]['machine_type'] == conso_energie.loc[j]['equipement_pilote_type_id']:
                        cumul_enr += coeffs_discontinu.loc[i]['count'] * conso_energie.loc[j]['consommation']
            for i in coeffs_continu.index:
                for j in conso_energie.index:
                    if coeffs_continu.loc[i]['machine_type'] == conso_energie.loc[j]['equipement_pilote_type_id']:
                        cumul_enr += coeffs_continu.loc[i]['count'] * conso_energie.loc[j]['consommation']
            return int(cumul_enr*1000)

    #Calcul du pourcentage de la conso d'énergie des foyers placée
    def conso_enr_placee() -> int:
        #On commence par calculer la quantité d'énergie placée dans les dernières 24h
        with conn_sortie.engine.connect() as conn:
            #On fabrique deux tables sous la forme machine_type | nombre de lancements
            cumul_enr_placee_24h:pd.DataFrame = pd.read_sql(""" SELECT SUM((0.25)*(p_c_with_flexible_consumption.power - p_c_without_flexible_consumption.power)) FROM p_c_with_flexible_consumption
                                                        INNER JOIN p_c_without_flexible_consumption
                                                        USING(data_timestamp)
                                                        WHERE p_c_with_flexible_consumption.data_timestamp > (CAST(EXTRACT (epoch FROM NOW()) AS INT)) """
                            , con = conn)
            cumul_enr_placee_24h = -cumul_enr_placee_24h
        cumul_enr_placee_24h = int(cumul_enr_placee_24h.loc[0]['sum'])
        print("cumul energie flexible",cumul_enr_placee_24h)
        #Calcul de la consommation d'énergie du panel mis à l'échelle les dernières 24h (item Panel_R_puissance_mae dans Zabbix)
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 60 * 24)
        puissance_panel_mae = 0
        for i in zapi.history.get(hostids = [10084], itemids = [44968], time_from = tf, time_till = tt, output = "extend", limit = 1440, history = 0):
            puissance_panel_mae += int(float(i['value'])) * (1/20) #Panel_R_puissance_mae
        pourcentage_enr_conso_placee = int(100*(cumul_enr_placee_24h/puissance_panel_mae))
        return pourcentage_enr_conso_placee

    def cumul_enr_autoconso() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = 1667287596
        #Calcul de la production mise à l'échelle sur l'heure (en Wh)
        enr_prod_mae = 0
        for i in zapi.history.get(hostids = [10084], itemids = [44969], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            enr_prod_mae += int(float(i['value']))*(1/20) #Panel_Prod_puissance_mae : 1 valeur toutes les 3 min
        #Calcul de la production en surplus à partir de l'équilibre (en Wh)
        surplus_prod = 0
        for i in zapi.history.get(hostids = [10084], itemids = [42883], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            if (float(i['value'])>0):
                surplus_prod += int(float(i['value']))*(1/60) #equilibre_general_p_c : 1 valeur par minute
        #Calcul de l'enr produite et consommée sur le territoire
        cumul_autoconso = int(enr_prod_mae - surplus_prod)
        return cumul_autoconso

    def cumul_enr_autoconso_opti() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Récupération de la précédente valeur
        val_prec = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45255], time_from = tf, time_till = tt-5, output = "extend", limit = 1, history=0):
            val_prec = int(float(i['value']))
        #Calcul de la production mise à l'échelle sur l'heure (en Wh)
        enr_prod_mae = 0
        for i in zapi.history.get(hostids = [10084], itemids = [44969], time_from = tf, time_till = tt, output = "extend", limit = 5, history=0):
            enr_prod_mae += int(float(i['value']))*(1/20) #Panel_Prod_puissance_mae : 1 valeur toutes les 3 min
        #Calcul de la production en surplus à partir de l'équilibre (en Wh)
        surplus_prod = 0
        for i in zapi.history.get(hostids = [10084], itemids = [42883], time_from = tf, time_till = tt, output = "extend", limit = 15, history=0):
            if (float(i['value'])>0):
                surplus_prod += int(float(i['value']))*(1/60) #equilibre_general_p_c : 1 valeur par minute
        #Calcul de l'enr produite et consommée sur le territoire
        cumul_supp = int(enr_prod_mae - surplus_prod)
        cumul_autoconso = val_prec + cumul_supp
        return cumul_autoconso

    def pourcentage_autoconso_mois() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 60 * 24 * datetime.now().day)
        #Calcul de l'enr produite et consommée sur le territoire pendant le mois courant
        enr_prod_mae = 0 #Production mise à l'échelle sur l'heure (en Wh)
        for i in zapi.history.get(hostids = [10084], itemids = [44969], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            enr_prod_mae += int(float(i['value']))*(1/20)
        surplus_prod = 0 #Production en surplus à partir de l'équilibre (en Wh)
        for i in zapi.history.get(hostids = [10084], itemids = [42883], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            if (float(i['value'])>0):
                surplus_prod += int(float(i['value']))*(1/60)
        enr_prod_et_conso = int(enr_prod_mae - surplus_prod) #Enr produite et consommée sur le territoire
        #Calcul de la production mise à l'échelle du panel pendant le mois courant
        enr_prod = 0
        for i in zapi.history.get(hostids = [10084], itemids = [44969], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            enr_prod += int(float(i['value'])) #Panel_Prod_puissance_mae
        #Calcul du pourcentage d'autoconsommation
        pourcentage_autoconso = int(100 * (enr_prod_et_conso/enr_prod))
        return pourcentage_autoconso

    def enr_eolien() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Calcul de l'énergie éolienne des 15 dernières minutes à partir de la puissance
        enr_eol = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45197], time_from = tf, time_till = tt, output = "extend", limit = 15, history=0):
            enr_eol += int(float(i['value']))*(1/60) #Prod_eolienne
        return int(enr_eol)

    def enr_solaire() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Calcul de l'énergie solaire des 15 dernières minutes à partir de la puissance
        enr_sol = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45198], time_from = tf, time_till = tt, output = "extend", limit = 15, history=0):
            enr_sol += int(float(i['value']))*(1/60) #Prod_solaire
        return int(enr_sol)

    def enr_metha() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Calcul de l'énergie métha des 15 dernières minutes à partir de la puissance
        enr_meth = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45248], time_from = tf, time_till = tt, output = "extend", limit = 15, history=0):
            enr_meth += int(float(i['value']))*(1/60) #Prod_methanisation
        return int(enr_meth)

    def part_eolien_prod_15min() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Production totale d'enr sur le dernier 1/4h
        prod_enr = 0
        #Production d'éolien sur le dernier 1/4h
        puissance_eolien = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45198], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_solaire
        for i in zapi.history.get(hostids = [10084], itemids = [45248], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_methanisation
        for i in zapi.history.get(hostids = [10084], itemids = [45197], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_eolienne
            puissance_eolien += int(i['value'])
        #Calcul du pourcentage
        pourcentage_prod_eolien = 100*puissance_eolien/prod_enr
        return int(pourcentage_prod_eolien)

    def part_solaire_prod_15min() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Production totale d'enr sur le dernier 1/4h
        prod_enr = 0
        #Production de solaire sur le dernier 1/4h
        puissance_solaire = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45198], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_solaire
            puissance_solaire += int(i['value'])
        for i in zapi.history.get(hostids = [10084], itemids = [45248], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_methanisation
        for i in zapi.history.get(hostids = [10084], itemids = [45197], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_eolienne
        #Calcul du pourcentage
        pourcentage_prod_solaire = 100*puissance_solaire/prod_enr
        return int(pourcentage_prod_solaire)

    def part_metha_prod_15min() -> int:
        #Connection au Zabbix
        zapi = createZapi()
        tt = int(time.mktime(datetime.now().timetuple()))
        tf = int(tt - 60 * 15)
        #Production totale d'enr sur le dernier 1/4h
        prod_enr = 0
        #Production de metha sur le dernier 1/4h
        puissance_metha = 0
        for i in zapi.history.get(hostids = [10084], itemids = [45198], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_solaire
        for i in zapi.history.get(hostids = [10084], itemids = [45248], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_methanisation
            puissance_metha += int(i['value'])
        for i in zapi.history.get(hostids = [10084], itemids = [45197], time_from = tf, time_till = tt, output = "extend", limit = 1440, history=0):
            prod_enr += int(i['value']) #Prod_eolienne
        #Calcul du pourcentage
        pourcentage_prod_metha = 100*puissance_metha/prod_enr
        return int(pourcentage_prod_metha)

    print("\n DEBUT")
    print(pourcentage_autoconso_mois())
    print("FIN \n")

    #Encapsulation dans un csv
    filename = "indics.csv"
    finalpath = os.path.join(path, filename)
    print("Path : ", finalpath)
    fichier = open("/home/indicateurs/Indicateurs-ELFE/indicateurs/indics.csv","w")

    res1 = str(indic_appareils_connectes())
    fichier.write("\"Zabbix server\" Nombre_appareils_connectes_test " + res1 + "\n")

    res2 = str(pourcentage_app_lancés_24h())
    fichier.write("\"Zabbix server\" Pourcentage_app_lances_24h_test " + res2 + "\n")

    res_cumul = str(cumul_enr())
    fichier.write("\"Zabbix server\" Cumul_energie_placee_test " + res_cumul + "\n")

    res_conso = str(conso_enr_placee())
    fichier.write("\"Zabbix server\" Pourcentage_energie_consommee_placee_test " + res_conso + "\n")

    res_cautoconso = str(cumul_enr_autoconso_opti())
    fichier.write("\"Zabbix server\" Energie_autoconsommee_test " + res_cautoconso + "\n")

    res_pautoconso = str(pourcentage_autoconso_mois())
    fichier.write("\"Zabbix server\" Pourcentage_autoconsommation_test " + res_pautoconso + "\n")

    res_enreol = str(enr_eolien())
    fichier.write("\"Zabbix server\" Enr_eolienne_test " + res_enreol + "\n")

    res_enrsol = str(enr_solaire())
    fichier.write("\"Zabbix server\" Enr_solaire_test " + res_enrsol + "\n")

    res_enrmeth = str(enr_metha())
    fichier.write("\"Zabbix server\" Enr_methanisation_test " + res_enrmeth + "\n")

    res_eol = str(part_eolien_prod_15min())
    fichier.write("\"Zabbix server\" Part_eolien_prod " + res_eol + "\n")

    res_sol = str(part_solaire_prod_15min())
    fichier.write("\"Zabbix server\" Part_solaire_prod " + res_sol + "\n")

    res_meth = str(part_metha_prod_15min())
    fichier.write("\"Zabbix server\" Part_metha_prod " + res_meth + "\n")

    #Connection au Zabbix
    try:
        zab = ConnectionZabbix('192.168.30.111', 'Zabbix server')
        print(f"Connection to Zabbix made successfully.")
    except Exception as ex:
        print("Connection to Zabbix could not be made due to the following error: \n", ex)

    def addMeasurement(cle:str, res:str):
        m = zb.Measurement(zab.host, cle, res)
        zab.measurements.add_measurement(m)

    addMeasurement("Pourcentage_app_lances_24h_test", res2)
    addMeasurement("Nombre_appareils_connectes_test", res1)
    addMeasurement("Cumul_energie_placee_test", res_cumul)
    addMeasurement("Pourcentage_energie_consommee_placee_test", res_conso)
    addMeasurement("Energie_autoconsommee_test", res_cautoconso)
    addMeasurement("Pourcentage_autoconsommation_test", res_pautoconso)
    addMeasurement("Enr_eolienne_test", res_enreol)
    addMeasurement("Enr_solaire_test", res_enrsol)
    addMeasurement("Enr_methanisation_test", res_enrmeth)

    try:
        await zab.response()
        print(f"Measurements well send to the Zabbix server")
    except Exception as ex:
        print("Measurements could not be send to the Zabbix server due to the following error: \n", ex)

if __name__ == "__main__":
    asyncio.run(main())


import pandas as pd
import requests
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib import transforms
import numpy as np

# Fetch data from the Elering API
start = "2022-04-01T00%3A00%3A00.000Z"
end = "2023-03-31T23%3A59%3A59.999Z"
url = f"https://dashboard.elering.ee/api/nps/price?start={start}&end={end}"
response = requests.get(url)

if response.status_code == 200:
    data = response.json()["data"]["ee"]
else:
    raise Exception(f"Failed to fetch data from Elering API. Status code: {response.status_code}")

# Convert data to a DataFrame
df = pd.DataFrame(data)
df["timestamp"] = pd.to_datetime(df["timestamp"], unit="s")
df["date"] = df["timestamp"].dt.date

# Convert fixed price from 15 eurocents/kWh to EUR/MWh (METHOD 1)
fixed_price = 15 * 10

# Calculate the average stock market electricity price for each separate day (METHOD 2)
daily_average_market_prices = df.groupby("date")["price"].mean()

# Calculate the cost using the 4 cheapest hours of each day (METHOD 3)
cheapest_hours_df = df.groupby("date", as_index=False).apply(lambda x: x.nsmallest(4, "price"))
cheapest_hours_daily_avg = cheapest_hours_df.groupby("date")["price"].mean()

# Calculate daily costs in EUR for consuming electricity for 4 hours every day, at 10 kW per hour (SCENARIO)
daily_consumption = 4  # in hours
daily_power_consumption_kw = 2  # in kilowatts (kW)
daily_energy_consumption_kwh = daily_power_consumption_kw * daily_consumption  # in kilowatt-hours (kWh)

# Calculate daily costs
fixed_daily_cost = fixed_price * daily_energy_consumption_kwh / 1000
average_market_daily_cost = daily_average_market_prices * daily_energy_consumption_kwh / 1000
cheapest_hours_daily_cost = cheapest_hours_daily_avg * daily_energy_consumption_kwh / 1000

# Calculate cumulative costs
fixed_cumulative_cost = np.cumsum(np.repeat(fixed_daily_cost, len(average_market_daily_cost)))
average_market_cumulative_cost = np.cumsum(average_market_daily_cost)
cheapest_hours_cumulative_cost = np.cumsum(cheapest_hours_daily_cost)

# CREATE A LINE GRAPH

# Set up the line graph
dates = daily_average_market_prices.index
plt.plot(dates, fixed_cumulative_cost, label="Fikseeritud hind (15 eurosenti/kWh)", linewidth=2)
plt.plot(dates, average_market_cumulative_cost, label="Päeva keskmine turuhind", linewidth=2)
plt.plot(dates, cheapest_hours_cumulative_cost, label="Päeva 4 odavamat tundi", linewidth=2)

# Add labels, title and legend
plt.ylabel("Summaarne kulu (EUR)", fontsize=14)
plt.xlabel("Kuupäev", fontsize=14)
plt.title("Eesti elektrikulude võrdlus (2022-04-01 kuni 2023-03-31)", fontsize=16)
plt.legend(fontsize=12)

# Configure x-axis ticks
ax = plt.gca()
ax.xaxis.set_major_locator(mdates.MonthLocator(interval=1))
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
plt.xticks(rotation=45, fontsize=12)
plt.yticks(fontsize=12)

# Add a textbox with the assumption
textbox_props = dict(boxstyle='round, pad=0.3', edgecolor='black', facecolor='lightgrey', linewidth=1, alpha=0.8)
ax = plt.gca()
transform = ax.transAxes
ax.text(0.01, 0.84, f"Eeldus: 4 tundi päevas, 2 kWh juures (kokku 8 kWh).", fontsize=12, bbox=textbox_props, transform=transform)

# Add gridlines
plt.grid()

# Add labels for last data points
plt.annotate("{:.2f} EUR".format(fixed_cumulative_cost[-1]), xy=(dates[-1], fixed_cumulative_cost[-1]), xytext=(-60, 8), textcoords='offset points', fontsize=14)
plt.annotate("{:.2f} EUR".format(average_market_cumulative_cost[-1]), xy=(dates[-1], average_market_cumulative_cost[-1]), xytext=(-60, 8), textcoords='offset points', fontsize=14)
plt.annotate("{:.2f} EUR".format(cheapest_hours_cumulative_cost[-1]), xy=(dates[-1], cheapest_hours_cumulative_cost[-1]), xytext=(-60, 8), textcoords='offset points', fontsize=14)

# Save it as "electricity_costs.png"
plt.savefig("electricity_costs.png", dpi=300)

# Display the graph
plt.show()


import base64
import hashlib
import json
from urllib import quote_plus

from internetofmoney.managers.BaseManager import BaseManager

class PayPalBaseManager(BaseManager):

    def __init__(self, database, cache_dir='cache'):
        super(PayPalBaseManager, self).__init__(database, cache_dir=cache_dir)
        self.persistent_storage['registered'] = True  # No registration required, one can immediately login
        self.logged_in = False
        self.access_token = None
        self.guid = "FB04564F-6675-454A-972E-9A50A5D75A8B"

        # Keep a cache with the descriptions that we've already seen for each transaction. This saves requests.
        self.known_descriptions = {}

    def is_logged_in(self):
        return self.logged_in

    def persistent_storage_filename(self):
        return 'paypal_account.json'

    def get_bank_name(self):
        return 'PayPal'

    def get_bank_id(self):
        return 'PP'

    def is_switch_capable(self):
        """
        Currently, PayPal is not capable of being a money switch.
        """
        return False

    def generate_auth_nonce(self, timestamp, email, password):
        data = "com.yourcompany.PPClient" + self.guid + self.guid + email + password[:3] + str(timestamp)
        sha_mod = hashlib.sha256()
        sha_mod.update(data)
        return base64.encodestring(sha_mod.digest()).rstrip('\n')

    def get_paypal_headers(self, is_json=False):
        if self.logged_in:
            auth_header = 'Bearer %s' % self.access_token
        else:
            auth_header = 'Basic QVY4aGRCQk04MHhsZ0tzRC1PYU9ReGVlSFhKbFpsYUN2WFdnVnB2VXFaTVRkVFh5OXBtZkVYdEUxbENxOg=='

        if is_json:
            content_type = 'application/json'
        else:
            content_type = 'application/x-www-form-urlencoded'

        app_context = {
            "deviceLocationCountry": "Unknown",
            "deviceOS": "iPhone OS",
            "appName": "com.yourcompany.PPClient",
            "sdkVersion": "1.0.0",
            "appGuid": self.guid,
            "appVersion": "6.11.0",
            "deviceLocale": "nl_NL",
            "deviceNetworkType": "Unknown",
            "deviceOSVersion": "8.1.2",
            "deviceLanguage": "nl",
            "deviceType": "iPhone",
            "deviceMake": "Apple",
            "deviceModel": "iPhone",
            "deviceNetworkCarrier": "Unknown"
        }

        return {
            'Accept-Language': 'nl-nl',
            'Accept': '*/*',
            'Connection': 'keep-alive',
            'Content-Type': content_type,
            'User-Agent': 'PayPal/116 (iPhone; iOS 8.1.2; Scale/2.00)',
            'X-PayPal-ConsumerApp-Context': quote_plus(json.dumps(app_context)),
            'Authorization': auth_header,
        }


# -*- coding: utf-8 -*-
"""Realise an optimisation of the hyper-parameters.
"""

__authors__ = "emenager, tnavez"
__contact__ = "etienne.menager@inria.fr, tanguy.navez@inria.fr"
__version__ = "1.0.0"
__copyright__ = "(c) 2022, Inria"
__date__ = "Nov 7 2022"

import sys
import os
import optuna
import pathlib
import torch
import torch.optim as optim
import numpy as np
import joblib

sys.path.insert(0, str(pathlib.Path(__file__).parent.absolute())+"/../")
sys.path.insert(0, str(pathlib.Path(__file__).parent.absolute()))

from Libraries.utils import ask_user_sampling_stat_input, compute_std_normalization, compute_minmax_normalization
from Libraries.Learning.DataSetManager import DataSetManager, get_data_loader
from Libraries.database import *
import Libraries.Learning.MLP.learning_tools as MLPlearning_tools

PATH = str(pathlib.Path(str(pathlib.Path(__file__).parent.absolute())))
os.makedirs(PATH + "/optimisation", exist_ok=True)

def init_optimisation(config, ratio_test_train = 0.25):
    """
    Init the dataset for the optimisation.

    Parameters
    ----------
    config: Config
        Config instance for the simulated scene
    ratio_test_train: float
    	The ratio between the train set and the test set.

    Outputs
    -------
    train_dataset: DataSet
        The train dataset, containing data for the learning step.
    test_dataset: DataSet
        The test dataset, containing data for the validation step.

    """
    n_samples, sampling_strategy, id_sampling = ask_user_sampling_stat_input(config)

    q_train = list(query_sampling_stats(model_name=config.model_name, sampling_strategy=sampling_strategy, n_samples=n_samples))
    n_test_samples = int(ratio_test_train * q_train[0]['n_samples'])
    q_test = list(query_sampling_stats(model_name = config.model_name, sampling_strategy = "Random", n_samples = n_test_samples))

    train_dataset = DataSetManager(config.model_name, id_sampling_stats = q_train[0]['id'], train = True)
    test_dataset = DataSetManager(config.model_name, id_sampling_stats = q_test[0]['id'], train = False)

    return train_dataset, test_dataset

def init_network(parameters, train_dataset, test_dataset, config):
    """
    Init the network for the optimisation.

    Parameters
    ----------
    parameters: dictionary
    	The parameters we optimize (batchsize, hiddensize, hiddenlayer, learning_rate).
    train_dataset, test_dataset: DataSet
        Dataset containing data for the learning step and the validation step.
    config: Config
        Config instance for the simulated scene

    Outputs
    -------
    dataloader, dataloader_test: torch dataloaders
        Dataloader for the learning step and the validation step.
    model: torch model
        The neural network we want to optimize.
    optimizer: torch.optim optimizer
        The optimizer for the neural network.
    scheduler: torch.optim.lr_scheduler.ReduceLROnPlateau
        The scheduler to change the value of the learning rate over time.
    """

    hiddensize = int(parameters["hiddensize"])
    hiddenlayer = int(parameters["hiddenlayer"])

    n_constraint = len(train_dataset.X[0][2]) # This may be changing when including collisions
    dataloader = get_data_loader(train_dataset, config.config_network["batch_size"], True)
    dataloader_test = get_data_loader(test_dataset, config.config_network["batch_size"], True) # Only one batch for test set

    LATENT_SIZE = hiddensize
    N_HIDDEN_LAYER = hiddenlayer

    OUTPUT_SIZE = len(np.triu_indices(n=n_constraint)[0]) + n_constraint
    INPUT_SIZE = OUTPUT_SIZE + n_constraint

    model = MLPlearning_tools.MLP(INPUT_SIZE, OUTPUT_SIZE, LATENT_SIZE, n_hidden_layers=N_HIDDEN_LAYER, dropout_probability=config.config_network["dropout_probability"])

    optimizer = optim.Adam(model.parameters(), lr=1e-3)  # Optimizer
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-05)

    return dataloader, dataloader_test, model, optimizer, scheduler

def eval_loop(model, optimizer, dataloader, scaling, config, is_train, previous_nb_update):
    """
        Test network against all dataloader set and compute loss.

        Parameters
        ----------
        model: torch model
            The neural network we want to optimize.
        optimizer: torch.optim optimizer
            The optimizer for the neural network.
        dataloader: torch dataloader
            Dataloader for the learning step/the validation step.
        scaling: list of tensor
            Scaling value for the data normalisation/standardisation
        config: Config
            Config instance for the simulated scene
        is_train: bool
            Train or test set
        previous_nb_update: int
            Number of gradient descent already done (= number of update)

        Outputs
        ----------
        e_loss: float
            Loss for evaluated model
        nb_update: int
            Number of update after the eval loop

        """

    e_loss = 0
    n_batches = 0
    nb_update = previous_nb_update

    for i, data in enumerate(dataloader):
        X_batch, Y_batch = data

        if config.config_network["data_normalization"] == "Std":
            X, Y = MLPlearning_tools.create_data_std(X_batch, Y_batch, scaling[0], scaling[1], scaling[2], scaling[3], design_to_MM=False)
        elif config.config_network["data_normalization"] == "MinMax":
            X, Y = MLPlearning_tools.create_data_minmax(X_batch, Y_batch, scaling[0], scaling[1], scaling[2],  scaling[3], design_to_MM=False)
        else:
            X, Y = MLPlearning_tools.create_data(X_batch, Y_batch, design_to_MM=False)

        Y_pred = model(X)
        if config.config_network["mode_loss"] != "physics_informed":
            loss = MLPlearning_tools.create_loss(Y_pred, Y, config.config_network["mode_loss"])
        else:
            loss = MLPlearning_tools.create_physics_informed_loss(Y_pred, Y, X, n_constraint, len(config.get_actuators_variables()),
                                    config.config_network["data_normalization"], scaling, is_contact=(len(config.get_contacts_variables()) != 0))

        if is_train:
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            nb_update+= 1

        e_loss += loss
        n_batches += 1
    e_loss /= n_batches

    return e_loss, nb_update

def main_loop(parameters, train_dataset, test_dataset, scaling, config):
    """
    Train the network using acquire data for a given set of parameters.

    Parameters
    ----------
    parameters: dictionary
    	The parameters we optimize (batchsize, hiddensize, hiddenlayer, learning_rate).
    train_dataset, test_dataset: DataSet
        Dataset containing data for the learning step and the validation step.
    scaling: list of tensor
            Scaling value for the data normalisation/standardisation
    config: Config
            Config instance for the simulated scene

    Outputs:
    -------
        test_loss: float
            The loss obtained on the validation set after training the network during n_epochs epochs.
    """
    dataloader, dataloader_test, model, optimizer, scheduler = init_network(parameters, train_dataset, test_dataset, config)
    nb_update, nb_update_max = 0, 1000
    while nb_update < nb_update_max:
        print(">> Number of update: {}/{}".format(nb_update, nb_update_max))
        model.train()
        e_loss, nb_update = eval_loop(model, optimizer, dataloader, scaling, config, True, nb_update)
        scheduler.step(e_loss)

    model.eval()
    with torch.no_grad():
        test_loss, nb_update = eval_loop(model, optimizer, dataloader_test, scaling, config, False, nb_update)

    return test_loss.item()

def objective(trial, train_dataset, test_dataset, scaling, config):
    """
    Objective function for the optimisation step.

    Parameters
    ----------
        trial: optuna trial
            The trial for the study.
        train_dataset, test_dataset: DataSet
            Dataset containing data for the learning step and the validation step.
        scaling: list of tensor
            Scaling value for the data normalisation/standardisation
        config: Config
                Config instance for the simulated scene

    Outputs:
    -------
        The evaluation of the network given a set of parameter.
    """
    params = { "hiddensize": trial.suggest_int("hiddensize", 32, 512, step = 1),
               "hiddenlayer": trial.suggest_int("hiddenlayer", 2, 4, step = 1)
            }

    return  main_loop(params, train_dataset, test_dataset, scaling, config)

def hyperparameters_optimisation(config, n_cores = 1, load = False, n_optimisation = 1000):
    """
    Function for the optimisation of the hyperparameters of the network.

    Parameters:
    ----------
    config: Config
        Config instance for the simulated scene
    n_cores: int
        The number of cores for parallel optimisation.
    load: bool
        Load the study to see results or create and run a new one.
    n_optimisation: int
        Number of optimisation steps (corresponding to n_optimisation*20*n_cores trials).
    """
    normalization_method = config.config_network["data_normalization"]

    if not load:
        print(">>  Init dataset ...")
        train_dataset, test_dataset = init_optimisation(config)

        scaling = []
        if normalization_method == "Std":
            mean_features_X, std_features_X, mean_features_Y, std_features_Y = compute_std_normalization(train_dataset,
                                                                                design_to_MM=False)
            scaling = [mean_features_X, std_features_X, mean_features_Y, std_features_Y]
        elif normalization_method == "MinMax":
            min_features_X, max_features_X, min_features_Y, max_features_Y = compute_minmax_normalization(train_dataset,
                                                                                design_to_MM=False)
            scaling = [min_features_X, max_features_X, min_features_Y, max_features_Y]

        print(">>  Start optimisation ...")
        study = optuna.create_study(direction='minimize')
        for i in range(n_optimisation):
            study.optimize(lambda trial: objective(trial, train_dataset, test_dataset, scaling, config), n_jobs=n_cores, n_trials=20*n_cores, timeout=None)
            joblib.dump(study, PATH + "/optimisation/study.pkl")
            print(">>  Best Params ", study.best_params)

    else:
        study=joblib.load(PATH + "/optimisation/study.pkl")

        print(">>  Best trial until now:")
        print(">>    Value: ", study.best_trial.value)
        print(">>    Params: ")

        params = dict(study.best_trial.params.items())
        print("hiddensize: {},".format(params["hiddensize"]))
        print("hiddenlayer: {},".format(params["hiddenlayer"]))


import math
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
from transformers import BertModel
d_model = 768
from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)

def gelu(x):
    """Implementation of the gelu activation function.
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

class BertLayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
        """Construct a layernorm module in the TF style (epsilon inside the square root).
        """
        super(BertLayerNorm, self).__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))    #
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x):
        u = x.mean(-1, keepdim=True)
        s = (x - u).pow(2).mean(-1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
        return self.weight * x + self.bias

class BertOutput(nn.Module):
    def __init__(self):
        super(BertOutput, self).__init__()
        self.dense = nn.Linear(3072, 768)
        self.LayerNorm = BertLayerNorm(768, eps=1e-12)
        self.dropout = nn.Dropout(0.1)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states

class BertIntermediate(nn.Module):
    def __init__(self):
        super(BertIntermediate, self).__init__()
        self.dense = nn.Linear(768, 3072)

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = gelu(hidden_states)
        return hidden_states

class BertSelfOutput(nn.Module):
    def __init__(self):
        super(BertSelfOutput, self).__init__()
        self.dense = nn.Linear(768, 768)
        self.LayerNorm = BertLayerNorm(768, eps=1e-12)
        self.dropout = nn.Dropout(0.1)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states

class KLBert(nn.Module):
    def __init__(self):
        super(KLBert, self).__init__()
        self.text_bert = BertModel.from_pretrained('bert-base-uncased', return_dict=False)
        self.know_bert = BertModel.from_pretrained('bert-base-uncased', return_dict=False)
        self.W_gate = nn.Linear(768 * 2, 1)
        self.intermediate = BertIntermediate()
        self.output = BertSelfOutput()
        self.classifier = nn.Linear(768, 2)
        self.secode_output = BertOutput()
        self.incongruence1 = Incongruence()
        self.incongruence2 = Incongruence()
        self.incongruence_cross = IncongruenceCross()

    def forward(self, text_ids, text_mask, know_ids, know_mask, labels=None):
        text_info, pooled_text_info = self.text_bert(input_ids=text_ids, attention_mask=text_mask)
        know_info, pooled_know_info = self.know_bert(input_ids=know_ids, attention_mask=know_mask)

        text_info_mask_mtr = get_attn_pad_mask(text_mask, text_mask)
        know_info_mask_mtr = get_attn_pad_mask(know_mask, know_mask)
        text_know_mask_mtr = get_attn_pad_mask(text_mask, know_mask)

        know_info_incong = self.incongruence2(know_info, know_info_mask_mtr,know_ids)
        text_info_incong = self.incongruence1(text_info, text_info_mask_mtr, text_ids)
        know_text_incong = self.incongruence_cross(text_info_incong, text_know_mask_mtr, know_info_incong, text_ids, know_ids)

        know_info_mean1 = torch.mean(know_info, dim=1).unsqueeze(1).expand(text_info.size(0),text_info.size(1),text_info.size(-1))
        text_info_mean1 = torch.mean(text_info, dim=1).unsqueeze(1).expand(text_info.size(0),text_info.size(1),text_info.size(-1))
        combine_info = torch.cat([text_info, know_info_mean1+text_info_mean1], dim=-1)

        # ablation study with "add" method
        '''res_add = torch.mean((text_info_incong + know_text_incong), dim=1, keepdim=True)
        res_add = self.add_lin2(self.add_lin1(res_add))
        intermediate_res = self.intermediate(res_add)
        # 32*1*768
        res = self.secode_output(intermediate_res, res_add)'''

        # ablation study with no "gate" method
        '''res_nogate = (text_info_incong[:, 0, :] + know_text_incong[:, 0, :]).unsqueeze(1)
        intermediate_res = self.intermediate(res_nogate)
        # 32*1*768
        res = self.secode_output(intermediate_res, res_nogate)'''

        alpha = self.W_gate(combine_info)
        alpha = torch.sigmoid(alpha)

        text_info_end= torch.matmul(alpha.transpose(1, 2), text_info_incong)
        know_text = torch.matmul((1 - alpha).transpose(1, 2), know_text_incong)

        res = self.output(know_text, text_info_end)

        intermediate_res = self.intermediate(res)
        res = self.secode_output(intermediate_res, res)
        logits = self.classifier(res)

        if labels is not None:

            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, 2), labels.view(-1))
            return loss
        else:
            return logits

class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.lin1 = nn.Linear(d_model, 3072)
        self.lin2 = nn.Linear(3072, d_model)
        self.layer_norm = BertLayerNorm(768, eps=1e-12)
        self.dropout = nn.Dropout(0.1)
    def forward(self, inputs):
        residual = inputs
        output = gelu(self.lin1(inputs))
        output = self.lin2(output)
        output = self.dropout(output)
        return self.layer_norm(output + residual)

def get_attn_pad_mask(seq_q, seq_k):
    batch_size, len_q= seq_q.size()
    batch_size, len_k= seq_k.size()
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)
    return pad_attn_mask.expand(batch_size, len_q, len_k)

class Incongruence(nn.Module):
    def __init__(self):
        super(Incongruence, self).__init__()
        self.incong_lin1 = nn.Linear(d_model, d_model)
        self.incong_lin2 = nn.Linear(d_model, d_model)
        self.incong_lin3 = nn.Linear(d_model, d_model)
        self.layer_norm = BertLayerNorm(768, eps=1e-12)
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, text_info, text_info_mask_mtr):
        text_info_incon1 = self.incong_lin1(text_info)
        text_info_incon2 = self.incong_lin2(text_info)
        text_info_incon3 = self.incong_lin3(text_info)
        text_info_cdist = torch.cdist(text_info_incon1, text_info_incon2)
        text_info_cdist1 = text_info_cdist.masked_fill(text_info_mask_mtr, -1e9)
        scores_dist = torch.softmax(text_info_cdist1, dim=-1)
        incon_info = torch.matmul(scores_dist, text_info_incon3)
        text_info_add = self.layer_norm(text_info + incon_info)
        text_info_add = self.pos_ffn(text_info_add)
        return text_info_add

class IncongruenceCross(nn.Module):
    def __init__(self):
        super(IncongruenceCross, self).__init__()
        self.incong_lin1 = nn.Linear(d_model, d_model)
        self.incong_lin2 = nn.Linear(d_model, d_model)
        self.incong_lin3 = nn.Linear(d_model, d_model)
        self.layer_norm = BertLayerNorm(768, eps=1e-12)
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, text_info, text_know_mask_mtr, know_info):
        text_info_incon_q = self.incong_lin1(text_info)
        know_info_incon_k = self.incong_lin2(know_info)
        know_info_incon_v = self.incong_lin3(know_info)
        text_info_cdist = torch.cdist(text_info_incon_q, know_info_incon_k)
        text_know_cdist1 = text_info_cdist.masked_fill(text_know_mask_mtr, -1e9)
        scores_dist = torch.softmax(text_know_cdist1, dim=-1)
        incon_info = torch.matmul(scores_dist, know_info_incon_v)
        text_info_add = self.layer_norm(text_info + incon_info)
        text_info_add = self.pos_ffn(text_info_add)
        return text_info_add


'''
Parses Yahoo Groups messages originally in JSON into e-mail readable format.
Must be placed in folder with JSON files to parse - however, could easily be
modified to work from single location.
'''

import json
import glob, os
import email
import html
import re

redacted = input("Should the archive be redacted [enter Y / N]? ")
seen = set()
threads = dict()
subjects = dict()

def redact(jsonMsg, jsonStr, mid):
    global msg

    #If chosen, redact identifying information from message headers to limit mass
    #data scraping.
    if redacted.lower() == "y" or "yes":
        msgcontent = email.message_from_string(jsonStr)
        #Use email module to parse and redact headers with possible identifying info
        #Note that content of headers is somewhat inconsistent from message to message,
        #so deletion is inclusive for all possible headers with identifying information.
        #Even if redaction is not chosen, headers do not in many cases contain
        #full e-mail addresses, but only the handle before "@".
        del msgcontent['from'], msgcontent['to'], msgcontent['cc'], msgcontent['return-path'], msgcontent['x-sender'], msgcontent['X-Yahoo-Profile']
        msgstr = "Post #%s:\n\n%s" % (mid, msgcontent)
        #Unescape HTML character codes
        msg = html.unescape(msgstr)
    else:
        msgcontent = email.message_from_string(jsonStr)
        msgstr = "Post #%s:\n\n%s" % (mid, msgcontent)
        #Unescape HTML character codes
        msg = html.unescape(msgstr)

    return msg

def threadparse(tid, mid, msg, threadSub):
    #Collects messages into threads, based on existing thread ID,
    #which is identical to the Message ID of the first message in
    #the thread. Also assigns each thread a subject for use in the
    #filename.
    if tid in seen:
        threads[tid].append(msg)
    else:
        threads[tid] = [msg]
        subjects[tid] = threadSub
    #Adds message to set of processed messages, so downthread messaages
    #can be matched with the initial post.
    seen.add(mid)

def parser():
    global jsonMsg
    global jsonStr
    global tid
    global mid
    global subject

    read_files = glob.glob("*.json")
    threads = dict()
    for f in read_files:
        with open(f, 'r', encoding="utf8") as current_file:
            #Collect data for functions
            raw = current_file.read()
            jsonMsg = json.loads(raw)
            #Collect Topic ID numbers
            tid = str(jsonMsg['ygData']['topicId'])
            #Collect Message ID numbers
            mid = str(jsonMsg['ygData']['msgId'])
            #Get raw email from JSON
            jsonStr = str(jsonMsg['ygData']['rawEmail'])
            #Get message subject
            if 'subject' in jsonMsg['ygData']:
                threadSub = str(jsonMsg['ygData']['subject'])
            else:
                threadSub = "No Subject"
            #Process text
            redact(jsonMsg, jsonStr, mid)
            threadparse(tid, mid, msg, threadSub)

    return jsonMsg
    return jsonStr
    return tid
    return mid
    return threadSub

def main():
    parser()
    #Write each thread to an individual, labeled textfile
    for key, value in threads.items():
        #Get and clean up subjects for file title
        if key in subjects:
            titleSub = str(subjects[key])
            #Shorten title to fit filename length limits
            titleShort = titleSub[:100] + (titleSub[100:] and '...')
            #Escape unusable characters from filename, like "?"
            title = re.sub('[^\w\-_\. ]', '_', titleShort)
        with open('Thread #%s - %s.txt' % (key, title), 'w', encoding='utf-8') as msgfile:
            #Write thread to textfile
            threadcount = len(value)
            msgfile.write("Number of Posts in Thread: %s\n\n" % threadcount)
            for item in value:
                if item in value[:-1]:
                    msgfile.write("%s\n\n-------Next Post in Thread-------\n\n" % item)
                else:
                    msgfile.write("%s" % item)

if __name__=="__main__":
   main()


from inspect import getsourcefile
import os.path as path, sys
current_dir = path.dirname(path.abspath(getsourcefile(lambda:0)))
sys.path.insert(0, current_dir[:current_dir.rfind(path.sep)])
from AP import *
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
from scipy import optimize
import matplotlib.ticker as ticker

#COLOR_STYLE = ["red","green","blue"]
Y_LABEL = r"Zählrate"
X_LABEL = r"Stromstärke in $mA$"
X_START =0 * 0.001
Y_START =0
X_END = 1.1 * 0.001
Y_END = 2200

X_MAJOR_TICK = 0.2 *0.001
Y_MAJOR_TICK =500
X_MINOR_TICK = 0.05 *0.001
Y_MINOR_TICK = 100
SAVE_AS = "./XST/Detektortotzeit.pdf"

path_ = "./XST/data.xls"

data = getTableFromCells("E5","O35",path_,"V5")
Totdat = [[i* 0.0001  for i in range(1,11)],[]]
Totdat[1] = [data[i][19] for i in range(1,11)]
print(Totdat)
#-------- curve fit
def func(x,r,t):
    return r * x *np.exp(-1*r*t*x)
def funcArr(x,arr):
    return func(x,arr[0],arr[1])

popt,perr = optimize.curve_fit(func,Totdat[0],Totdat[1])
print(popt,np.sqrt(np.diag(perr)))
fitDat =genDataFromFunktion(1000,X_START,X_END,popt,funcArr)

#----------- plot

fig, ax = plt.subplots()

ax.grid()

ax.scatter(Totdat[0],Totdat[1],s=15)
ax.errorbar(Totdat[0],Totdat[1],xerr=20e-6,fmt="x", ecolor = 'black',elinewidth=0.9,capsize=4,capthick=0.9,label="Messdaten")
ax.plot(fitDat[0],fitDat[1],color = "red")

ax.set_xlabel(X_LABEL)
ax.set_ylabel(Y_LABEL)
ax.set_xlim(X_START,X_END)
ax.set_ylim(Y_START,Y_END)

# Define the formatter function
def format_func(value, tick_number):
    return round(value*1000,2)

# Set the x-axis formatter
ax.xaxis.set_major_formatter(ticker.FuncFormatter(format_func))

ax.xaxis.set_major_locator(MultipleLocator(X_MAJOR_TICK))
ax.xaxis.set_minor_locator(MultipleLocator(X_MINOR_TICK))
ax.yaxis.set_major_locator(MultipleLocator(Y_MAJOR_TICK))
ax.yaxis.set_minor_locator(MultipleLocator(Y_MINOR_TICK))

plt.legend(("Messwerte", fr"fit mit r={round_errtex(popt[0],np.sqrt(np.diag(perr))[0])}; $\tau$ = {round_errtex(popt[1],np.sqrt(np.diag(perr))[1])}"))
plt.savefig(SAVE_AS)
plt.show()