import torch

import triton
import triton.language as tl

@triton.jit
def max_fn(x, y):
    return tl.math.max(x, y)

@triton.jit
def _fwd_kernel(
    Q, K, V, sm_scale,
    L,
    Out,
    stride_qz, stride_qh, stride_qm, stride_qk,
    stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vk, stride_vn,
    stride_oz, stride_oh, stride_om, stride_on,
    Z, H, N_CTX,
    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
    BLOCK_N: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
):
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    qvk_offset = off_hz * stride_qh
    Q_block_ptr = tl.make_block_ptr(
        base=Q + qvk_offset,
        shape=(N_CTX, BLOCK_DMODEL),
        strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL),
        order=(1, 0)
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + qvk_offset,
        shape=(BLOCK_DMODEL, N_CTX),
        strides=(stride_kk, stride_kn),
        offsets=(0, 0),
        block_shape=(BLOCK_DMODEL, BLOCK_N),
        order=(0, 1)
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + qvk_offset,
        shape=(N_CTX, BLOCK_DMODEL),
        strides=(stride_vk, stride_vn),
        offsets=(0, 0),
        block_shape=(BLOCK_N, BLOCK_DMODEL),
        order=(1, 0)
    )
    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)
    # initialize pointer to m and l
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    # scale sm_scale by log_2(e) and use
    # 2^x instead of exp in the loop because CSE and LICM
    # don't work as expected with `exp` in the loop
    qk_scale = sm_scale * 1.44269504
    # load q: it will stay in SRAM throughout
    q = tl.load(Q_block_ptr)
    q = (q * qk_scale).to(tl.float16)
    # loop over k, v and update accumulator
    lo = 0
    hi = (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX
    for start_n in range(lo, hi, BLOCK_N):
        # -- load k, v --
        k = tl.load(K_block_ptr)
        v = tl.load(V_block_ptr)
        # -- compute qk ---
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        if IS_CAUSAL:
            qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
        qk += tl.dot(q, k)
        # -- compute scaling constant ---
        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
        alpha = tl.math.exp2(m_i - m_i_new)
        p = tl.math.exp2(qk - m_i_new[:, None])
        # -- scale and update acc --
        acc_scale = l_i * 0 + alpha  # workaround some compiler bug
        acc *= acc_scale[:, None]
        acc += tl.dot(p.to(tl.float16), v)
        # -- update m_i and l_i --
        l_i = l_i * alpha + tl.sum(p, 1)
        m_i = m_i_new
        # update pointers
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
    # write back l and m
    acc = acc / l_i[:, None]
    l_ptrs = L + off_hz * N_CTX + offs_m
    tl.store(l_ptrs, m_i + tl.math.log2(l_i))
    # write back O
    O_block_ptr = tl.make_block_ptr(
        base=Out + qvk_offset,
        shape=(N_CTX, BLOCK_DMODEL),
        strides=(stride_om, stride_on),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL),
        order=(1, 0)
    )
    tl.store(O_block_ptr, acc.to(tl.float16))

@triton.jit
def _bwd_preprocess(
    Out, DO,
    Delta,
    BLOCK_M: tl.constexpr, D_HEAD: tl.constexpr,
):
    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
    off_n = tl.arange(0, D_HEAD)
    # load
    o = tl.load(Out + off_m[:, None] * D_HEAD + off_n[None, :]).to(tl.float32)
    do = tl.load(DO + off_m[:, None] * D_HEAD + off_n[None, :]).to(tl.float32)
    # compute
    delta = tl.sum(o * do, axis=1)
    # write-back
    tl.store(Delta + off_m, delta)

@triton.jit
def _bwd_kernel(
    Q, K, V, sm_scale, Out, DO,
    DQ, DK, DV,
    L,
    D,
    stride_qz, stride_qh, stride_qm, stride_qk,
    stride_kz, stride_kh, stride_kn, stride_kk,
    stride_vz, stride_vh, stride_vk, stride_vn,
    Z, H, N_CTX,
    num_block,
    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
    BLOCK_N: tl.constexpr,
    CAUSAL: tl.constexpr,
):
    off_hz = tl.program_id(0)
    off_z = off_hz // H
    off_h = off_hz % H
    qk_scale = sm_scale * 1.44269504
    # offset pointers for batch/head
    Q += off_z * stride_qz + off_h * stride_qh
    K += off_z * stride_qz + off_h * stride_qh
    V += off_z * stride_qz + off_h * stride_qh
    DO += off_z * stride_qz + off_h * stride_qh
    DQ += off_z * stride_qz + off_h * stride_qh
    DK += off_z * stride_qz + off_h * stride_qh
    DV += off_z * stride_qz + off_h * stride_qh
    for start_n in range(0, num_block):
        if CAUSAL:
            lo = start_n * BLOCK_M
        else:
            lo = 0
        # initialize row/col offsets
        offs_qm = lo + tl.arange(0, BLOCK_M)
        offs_n = start_n * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_m = tl.arange(0, BLOCK_N)
        offs_k = tl.arange(0, BLOCK_DMODEL)
        # initialize pointers to value-like data
        q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk)
        k_ptrs = K + (offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk)
        v_ptrs = V + (offs_n[:, None] * stride_qm + offs_k[None, :] * stride_qk)
        do_ptrs = DO + (offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk)
        dq_ptrs = DQ + (offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk)
        # pointer to row-wise quantities in value-like data
        D_ptrs = D + off_hz * N_CTX
        l_ptrs = L + off_hz * N_CTX
        # initialize dv amd dk
        dv = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
        dk = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
        # k and v stay in SRAM throughout
        k = tl.load(k_ptrs)
        v = tl.load(v_ptrs)
        # loop over rows
        for start_m in range(lo, num_block * BLOCK_M, BLOCK_M):
            offs_m_curr = start_m + offs_m
            # load q, k, v, do on-chip
            q = tl.load(q_ptrs)
            # recompute p = softmax(qk, dim=-1).T
            if CAUSAL:
                qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), float(0.), float("-inf"))
            else:
                qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
            qk += tl.dot(q, tl.trans(k))
            qk *= qk_scale
            l_i = tl.load(l_ptrs + offs_m_curr)
            p = tl.math.exp2(qk - l_i[:, None])
            # compute dv
            do = tl.load(do_ptrs)
            dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
            # compute dp = dot(v, do)
            Di = tl.load(D_ptrs + offs_m_curr)
            dp = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - Di[:, None]
            dp += tl.dot(do, tl.trans(v))
            # compute ds = p * (dp - delta[:, None])
            ds = p * dp * sm_scale
            # compute dk = dot(ds.T, q)
            dk += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q)
            # compute dq
            dq = tl.load(dq_ptrs)
            dq += tl.dot(ds.to(Q.dtype.element_ty), k)
            tl.store(dq_ptrs, dq)
            # increment pointers
            dq_ptrs += BLOCK_M * stride_qm
            q_ptrs += BLOCK_M * stride_qm
            do_ptrs += BLOCK_M * stride_qm
        # write-back
        dv_ptrs = DV + (offs_n[:, None] * stride_qm + offs_k[None, :] * stride_qk)
        dk_ptrs = DK + (offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk)
        tl.store(dv_ptrs, dv)
        tl.store(dk_ptrs, dk)

empty = torch.empty(128, device="cuda")

class _attention(torch.autograd.Function):

    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal, sm_scale):
        # shape constraints
        Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
        assert Lq == Lk and Lk == Lv
        assert Lk in {16, 32, 64, 128}
        o = torch.empty_like(q)
        BLOCK_M = 128
        BLOCK_N = 64
        grid = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1], 1)
        L = torch.empty((q.shape[0] * q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)

        num_warps = 4 if Lk <= 64 else 8
        _fwd_kernel[grid](
            q, k, v, sm_scale,
            L,
            o,
            q.stride(0), q.stride(1), q.stride(2), q.stride(3),
            k.stride(0), k.stride(1), k.stride(2), k.stride(3),
            v.stride(0), v.stride(1), v.stride(2), v.stride(3),
            o.stride(0), o.stride(1), o.stride(2), o.stride(3),
            q.shape[0], q.shape[1], q.shape[2],
            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=Lk,
            IS_CAUSAL=causal,
            num_warps=num_warps,
            num_stages=4)

        ctx.save_for_backward(q, k, v, o, L)
        ctx.grid = grid
        ctx.sm_scale = sm_scale
        ctx.BLOCK_DMODEL = Lk
        ctx.causal = causal
        return o

    @staticmethod
    def backward(ctx, do):
        BLOCK = 128
        q, k, v, o, L = ctx.saved_tensors
        do = do.contiguous()
        dq = torch.zeros_like(q, dtype=torch.float32)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)
        delta = torch.empty_like(L)
        _bwd_preprocess[(ctx.grid[0] * ctx.grid[1], )](
            o, do,
            delta,
            BLOCK_M=BLOCK, D_HEAD=ctx.BLOCK_DMODEL,
        )
        _bwd_kernel[(ctx.grid[1],)](
            q, k, v, ctx.sm_scale,
            o, do,
            dq, dk, dv,
            L, delta,
            q.stride(0), q.stride(1), q.stride(2), q.stride(3),
            k.stride(0), k.stride(1), k.stride(2), k.stride(3),
            v.stride(0), v.stride(1), v.stride(2), v.stride(3),
            q.shape[0], q.shape[1], q.shape[2],
            ctx.grid[0],
            BLOCK_M=BLOCK, BLOCK_N=BLOCK,
            BLOCK_DMODEL=ctx.BLOCK_DMODEL, num_warps=8,
            CAUSAL=ctx.causal,
            num_stages=1,
        )
        return dq, dk, dv, None, None

attention = _attention.apply


import datetime
import json

import pytest

from ..exchange_api import CurrencyExchange

def test_parse_conversion_rate():
    exchange = CurrencyExchange()
    path = './utils/tests/sample_data/exr-d-gbp-eur-20200224.json'
    with open(path, 'r') as fp:
        data = json.load(fp)
        rate = exchange._parse_exchange_rate(data)
        assert(str(rate) == '0.83833')

def test_exchange_rate_url():
    exchange = CurrencyExchange()
    actual = exchange._make_exchange_rate_url(currency='USD', date=datetime.date(2020, 2, 24))
    expected = 'https://sdw-wsrest.ecb.europa.eu/service/data/EXR/M.USD.EUR.SP00.A?startPeriod=2020-02-01&endPeriod=2020-02-29'
    assert(actual == expected)

def test_parse_last_modified():
    exchange = CurrencyExchange()
    headers = { 'Last-Modified': 'Mon, 24 Feb 2020 14:56:43 GMT' }
    actual = exchange._parse_last_modified(headers)
    expected = datetime.date(2020, 2, 24)
    assert(actual == expected)

def test_parse_invalid_last_modified():
    exchange = CurrencyExchange()
    headers = { 'Last-Modified': 'Gor, 37 Bar 9090 27:63:99 FEZ' }
    with pytest.raises(ValueError):
        exchange._parse_last_modified(headers)

def test_parse_no_last_modified():
    exchange = CurrencyExchange()
    headers = { 'Date': 'Mon, 24 Feb 2020 14:56:43 GMT' }
    with pytest.raises(ValueError):
        exchange._parse_last_modified(headers)

def test_get_exchange_rate():
    exchange = CurrencyExchange()
    rate = exchange.get_exchange_rate(currency='USD', date=datetime.date(2020, 2, 23))
    assert(str(rate) == '1.110036363636364')


import os
import numpy as np
from transformers import RobertaTokenizer, AutoTokenizer
from datasets import load_metric, load_from_disk, load_dataset, DatasetDict

base_model = "Pipper/SolCoder"
tokenizer = RobertaTokenizer.from_pretrained(base_model)

if not os.path.exists('./sol_dataset'):
    dataset = load_dataset("Pipper/sol_processed_s2s")
else:
    dataset = load_from_disk('./sol_dataset', keep_in_memory=True)

assert 'train' in dataset.keys()
assert 'valid' in dataset.keys()
assert 'test' in dataset.keys()

train_set = dataset['train']
eval_set = dataset['valid']
test_set = dataset['test']

# grap a random sample and print comment code pair
idx = 30
sample = train_set[idx]
print('Comment:', tokenizer.decode(sample['input_ids'], skip_special_tokens=True))
labels = np.where(np.array(sample['labels']) != -100, sample['labels'], tokenizer.pad_token_id)

print('code:\n', tokenizer.decode(labels, skip_special_tokens=True))

print('Verification succeed!')


import torch.nn as nn
from transformers import DistilBertModel

class DistilSequentialRecommender(nn.Module):
    def __init__(self, distilbert_model_name, sequence_length, hidden_size, num_classes):
        super(DistilSequentialRecommender, self).__init__()
        self.distilbert = DistilBertModel.from_pretrained(distilbert_model_name)
        self.pooling = nn.AvgPool1d(sequence_length)
        self.fc = nn.Linear(self.distilbert.config.dim, hidden_size)
        self.output_layer = nn.Linear(hidden_size, num_classes)

    def forward(self, input_ids, attention_mask):
        distilbert_output = self.distilbert(input_ids, attention_mask=attention_mask)[0]
        pooled_output = self.pooling(distilbert_output.permute(0, 2, 1)).squeeze(2)
        fc_output = self.fc(pooled_output)
        logits = self.output_layer(fc_output)
        return logits


from .config import (
    LRN_RATE,
    IMG_CHANNELS,
    LOAD_MODEL,
    SAVE_MODEL,
    DEVICE,
    NUM_EPOCHS,
    NUM_FEATURES,
    DEPTH_OF_CRITIC,
    MODEL_DIR,
    IMG_SIZE,
)
import torch
import torch.nn as nn
from .dataset import getLoaders
from .losses import CycleConsistencyLoss, WLossGP_Critic, WLoss_Generator, IdentityLoss
from .models import Generator, Critic
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import os

# train class for gans
from .utils import plot_to_tensorboard

class GANTrain(nn.Module):
    def __init__(self, device):
        super(GANTrain, self).__init__()
        self.device = device
        self.net_celeb_G = Generator(IMG_CHANNELS, IMG_CHANNELS).to(device)
        self.net_anime_G = Generator(IMG_CHANNELS, IMG_CHANNELS).to(device)
        self.net_celeb_C = Critic(IMG_CHANNELS, DEPTH_OF_CRITIC,
                                  NUM_FEATURES).to(device)
        self.net_anime_C = Critic(IMG_CHANNELS, DEPTH_OF_CRITIC,
                                  NUM_FEATURES).to(device)

        self.optim_celeb_G = optim.Adam(self.net_celeb_G.parameters(),
                                        lr=LRN_RATE)
        self.optim_anime_G = optim.Adam(self.net_anime_G.parameters(),
                                        lr=LRN_RATE)
        self.optim_celeb_C = optim.Adam(self.net_celeb_C.parameters(),
                                        lr=LRN_RATE)
        self.optim_anime_C = optim.Adam(self.net_anime_C.parameters(),
                                        lr=LRN_RATE)

        self.criterion_cycle_anime_celeb = CycleConsistencyLoss(
            self.net_anime_G, self.net_celeb_G)
        self.criterion_cycle_celeb_anime = CycleConsistencyLoss(
            self.net_celeb_G, self.net_anime_G)
        self.criterion_id_anime = IdentityLoss(self.net_anime_G)
        self.criterion_id_celeb = IdentityLoss(self.net_celeb_G)
        self.criterion_w_celeb_G = WLoss_Generator()
        self.criterion_w_anime_G = WLoss_Generator()
        self.criterion_w_celeb_C = WLossGP_Critic(self.net_celeb_C)
        self.criterion_w_anime_C = WLossGP_Critic(self.net_anime_C)

        self.anime_loader, self.celeb_loader = getLoaders(size=(IMG_SIZE,
                                                                IMG_SIZE))

        self.writer = SummaryWriter("logs/cycleGAN")
        self.tensorboard_step = 0

        if LOAD_MODEL:
            self.__load_checkpoint__()

        self.net_anime_C.train()
        self.net_celeb_C.train()
        self.net_celeb_G.train()
        self.net_celeb_G.train()

    def __save_checkpoint__(self):
        print("=>> Saving Checkpoint")
        self.net_celeb_G.save(self.optim_celeb_G,
                              name_file=os.path.join(MODEL_DIR,
                                                     'celeb_gen.pth.tar'))
        self.net_anime_G.save(self.optim_anime_G,
                              name_file=os.path.join(MODEL_DIR,
                                                     'anime_gen.pth.tar'))
        self.net_celeb_C.save(self.optim_celeb_C,
                              name_file=os.path.join(MODEL_DIR,
                                                     'celeb_cri.pth.tar'))
        self.net_anime_C.save(self.optim_anime_C,
                              name_file=os.path.join(MODEL_DIR,
                                                     'anime_cri.pth.tar'))

    def __load_checkpoint__(self):
        self.net_celeb_G.load(self.optim_celeb_G,
                              name_file=os.path.join(MODEL_DIR,
                                                     'celeb_gen.pth.tar'))
        self.net_anime_G.load(self.optim_anime_G,
                              name_file=os.path.join(MODEL_DIR,
                                                     'anime_gen.pth.tar'))
        self.net_celeb_C.load(self.optim_celeb_C,
                              name_file=os.path.join(MODEL_DIR,
                                                     'celeb_cri.pth.tar'))
        self.net_anime_C.load(self.optim_anime_C,
                              name_file=os.path.join(MODEL_DIR,
                                                     'anime_cri.pth.tar'))

    def __train_one__(self):
        loop = tqdm(zip(self.anime_loader, self.celeb_loader), leave=True)
        batch_idx = 0

        for real_anime, real_celeb in loop:
            self.optim_celeb_C.zero_grad()
            self.optim_anime_C.zero_grad()

            ## CRITIC LOSS ##
            real_celeb = real_celeb.to(self.device)
            real_anime = real_anime.to(self.device)

            # print(real_celeb.shape)
            # print(real_anime.shape)

            # FOR CELEBS
            fake_celeb = self.net_celeb_G(real_anime)
            critic_out_fake_celeb = self.net_celeb_C(fake_celeb.detach())
            critic_out_real_celeb = self.net_celeb_C(real_celeb)

            loss_celeb_critic = self.criterion_w_celeb_C(
                prediction_real=critic_out_real_celeb,
                prediction_fake=critic_out_fake_celeb,
                real=real_celeb,
                fake=fake_celeb,
            )
            ccloss_celeb = self.criterion_cycle_celeb_anime(real_celeb)
            idloss_celeb = self.criterion_id_celeb(real_celeb)

            sum_loss_celeb = loss_celeb_critic + ccloss_celeb + idloss_celeb

            sum_loss_celeb.backward(retain_graph=True)
            self.optim_celeb_C.step()

            # FOR ANIMES
            fake_anime = self.net_anime_G(real_celeb)
            critic_out_real_anime = self.net_anime_C(real_anime)
            critic_out_fake_anime = self.net_anime_C(fake_anime.detach())

            loss_anime_critic = self.criterion_w_anime_C(
                prediction_real=critic_out_real_anime,
                prediction_fake=critic_out_fake_anime,
                real=real_anime,
                fake=fake_anime,
            )

            ccloss_anime = self.criterion_cycle_anime_celeb(real_anime)
            idloss_anime = self.criterion_id_anime(real_anime)

            sum_loss_anime = loss_anime_critic + ccloss_anime + idloss_anime

            sum_loss_anime.backward(retain_graph=True)
            self.optim_anime_C.step()

            loss_critic = sum_loss_anime + sum_loss_anime

            ## GENERATOR LOSS ##

            # For Celebs
            self.optim_celeb_G.zero_grad()
            gen_celeb_fake = self.net_celeb_C(fake_celeb)
            loss_gen_celeb = self.criterion_w_celeb_G(gen_celeb_fake)

            loss_gen_celeb.backward()
            self.optim_celeb_G.step()

            # For Anime
            self.optim_anime_G.zero_grad()
            gen_anime_fake = self.net_anime_C(fake_anime)
            loss_gen_anime = self.criterion_w_anime_G(gen_anime_fake)

            loss_gen_anime.backward()
            self.optim_anime_G.step()

            loss_gen = loss_gen_anime + loss_gen_celeb

            if batch_idx % 5 == 0:
                with torch.no_grad():
                    fake_celeb = self.net_celeb_G(real_anime)
                    fake_anime = self.net_anime_G(real_celeb)
                    plot_to_tensorboard(
                        self.writer,
                        real_anime.detach(),
                        fake_celeb.detach(),
                        self.tensorboard_step,
                        loss_critic.item(),
                        loss_gen.item(),
                    )
                    plot_to_tensorboard(
                        self.writer,
                        real_celeb.detach(),
                        fake_anime.detach(),
                        self.tensorboard_step,
                        None,
                        None,
                    )

                    self.tensorboard_step += 1

            loop.set_postfix(loss_critic=loss_critic.item(),
                             loss_gen=loss_gen.item())
            batch_idx += 1

    def forward(self, num_epochs):
        for epoch in range(num_epochs):
            print(f"Epoch [{epoch + 1} / {num_epochs}]")
            self.__train_one__()

            if SAVE_MODEL:
                self.__save_checkpoint__()

def main():
    trainObj = GANTrain(DEVICE)
    trainObj(NUM_EPOCHS)

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
# ---------------------

from typing import *

import torch
import torch.nn as nn
from torch.nn.functional import interpolate

from models import BaseModel

class Upsample(nn.Module):
    """
    Upsamples a given tensor by (scale_factor)X.
    """

    def __init__(self, scale_factor=2, mode='trilinear'):
        # type: (int, str) -> Upsample
        """
        :param scale_factor: the multiplier for the image height / width
        :param mode: the upsampling algorithm - values in {'nearest', 'linear', 'bilinear', 'trilinear'}
        """
        super().__init__()
        self.scale_factor = scale_factor
        self.mode = mode

    def forward(self, x):
        # type: (torch.Tensor) -> torch.Tensor
        return interpolate(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=False)

    def extra_repr(self):
        return f'scale_factor={self.scale_factor}, mode={self.mode}'

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

class PixelNorm(nn.Module):
    """
    Normalize the feature vector in each pixel to unit length.
    """

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

    def forward(self, input):
        # type: (torch.Tensor) -> torch.Tensor
        return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)

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

class BasicBlock(nn.Module):

    def __init__(self, in_ch, out_ch, mode, pixel_norm, pad=2):
        # type: (int, int, str, bool, Union[int, Tuple[int, int]]) -> BasicBlock
        """
        :param in_ch: numebr of input channels
        :param out_ch: number of output channels
        :param mode: values in {'++', '--', '=='}
        :param pixel_norm: use pixel normalization?
        :param pad: padding
        """
        super().__init__()

        assert mode in ['++', '--', '==']

        self.block = nn.Sequential(
            nn.Conv3d(in_channels=in_ch, out_channels=out_ch, kernel_size=5, padding=pad),
            PixelNorm(), nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=out_ch, out_channels=out_ch, kernel_size=5, padding=pad),
            PixelNorm(), nn.ReLU(inplace=True),
        )

        if not pixel_norm:
            del self.block[1]
            del self.block[3]

        if mode in ['++', '--']:
            if mode == '++':
                final_block_element = Upsample(scale_factor=2)
            else:
                final_block_element = nn.MaxPool3d(kernel_size=2)
            self.block = nn.Sequential(self.block, final_block_element)

    def forward(self, x):
        # type: (torch.Tensor) -> torch.Tensor
        out = self.block(x)
        return out

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

class Autoencoder(BaseModel):

    def __init__(self, hmap_d):
        # type: (int) -> None
        """
        :param hmap_d: number of input channels
        """
        print(f'model: {__file__}\n')
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels=hmap_d, out_channels=hmap_d // 2, kernel_size=5, stride=2, padding=2),
            nn.ReLU(True),
            nn.Conv2d(in_channels=hmap_d // 2, out_channels=hmap_d // 4, kernel_size=5, stride=2, padding=2),
            nn.ReLU(True),
            nn.Conv2d(in_channels=hmap_d // 4, out_channels=hmap_d // 4, kernel_size=5, stride=1, padding=2),
            nn.ReLU(True),
        )

        self.fuser = nn.Sequential(
            nn.Conv3d(in_channels=14, out_channels=4, kernel_size=5, padding=2),
            nn.ReLU(True),
            nn.Conv3d(in_channels=4, out_channels=1, kernel_size=5, padding=2),
            nn.ReLU(True),
        )

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

        self.defuser = nn.Sequential(
            nn.Conv3d(in_channels=1, out_channels=4, kernel_size=5, padding=2),
            nn.ReLU(True),
            nn.Conv3d(in_channels=4, out_channels=14, kernel_size=5, padding=2),
            nn.ReLU(True)
        )

        self.decoder = nn.Sequential(
            nn.Conv2d(in_channels=hmap_d // 4, out_channels=hmap_d // 4, kernel_size=5, padding=2),
            Upsample(mode='bilinear'),
            nn.ReLU(True),
            nn.Conv2d(in_channels=hmap_d // 4, out_channels=hmap_d // 2, kernel_size=5, padding=2),
            Upsample(mode='bilinear'),
            nn.ReLU(True),
            nn.Conv2d(in_channels=hmap_d // 2, out_channels=hmap_d, kernel_size=5, padding=2),
            nn.ReLU(True)
        )

    def encode(self, x):
        # type: (torch.Tensor) -> torch.Tensor

        x = self.encoder(torch.reshape(x, (x.shape[0] * 14, x.shape[2], x.shape[3], x.shape[4])).contiguous())
        x = torch.reshape(x, (x.shape[0] // 14, 14, x.shape[1], x.shape[2], x.shape[3])).contiguous()

        x = self.fuser(x)
        return x

    def decode(self, x):
        x = self.defuser(x)

        x = self.decoder(torch.reshape(x, (x.shape[0] * 14, x.shape[2], x.shape[3], x.shape[4])).contiguous())
        x = torch.reshape(x, (x.shape[0] // 14, 14, x.shape[1], x.shape[2], x.shape[3])).contiguous()
        return x

    def forward(self, x):
        # type: (torch.Tensor) -> torch.Tensor
        x = self.encode(x)
        x = self.decode(x)
        return x

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

def main():
    import time
    import numpy as np
    from conf import Conf

    cnf = Conf(exp_name='default')

    model = Autoencoder(hmap_d=cnf.hmap_d).to(cnf.device)
    model.load_w('/home/fabio/PycharmProjects/LoCO/models/weights/vha.pth')

    print(model)
    print(f'* number of parameters: {model.n_param}')

    x = torch.rand((1, 14, cnf.hmap_d, cnf.hmap_h, cnf.hmap_w)).to(cnf.device)

    print('\n--- ENCODING ---')
    y = model.encode(x)
    print(f'* input shape: {tuple(x.shape)}')
    print(f'* output shape: {tuple(y.shape)}')

    print(f'* space savings: {100 - 100 * np.prod(tuple(y.shape)) / np.prod(tuple(x.shape)):.2f}%')
    print(f'* factor: {np.prod(tuple(x.shape)) / np.prod(tuple(y.shape)):.2f}')

    print('\n--- DECODING ---')
    xd = model.decode(y)
    print(f'* input shape: {tuple(y.shape)}')
    print(f'* output shape: {tuple(xd.shape)}')

    print('\n--- FORWARD ---')
    y = model.forward(x)
    print(f'* input shape: {tuple(x.shape)}')
    print(f'* output shape: {tuple(y.shape)}')

if __name__ == '__main__':
    main()


from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
from datetime import datetime
from selenium.common.exceptions import NoSuchElementException
import xlsxwriter

trafficWorkbook=xlsxwriter.Workbook("Traffic record.xlsx")
trafficSheet=trafficWorkbook.add_worksheet()

URL="https://www.google.com/maps/dir/1239+FL-436+%23101,+Casselberry,+FL+32707/2101+Water+Bridge+Blvd,+Orlando,+FL+32837/@28.5319064,-81.517875,11z/data=!3m1!4b1!4m13!4m12!1m5!1m1!1s0x88e76e36b22d92eb:0x4c749fb2c5492ea5!2m2!1d-81.3231343!2d28.6347265!1m5!1m1!1s0x88e77d8ac5cac91b:0x10c348dd5c0c709a!2m2!1d-81.4050814!2d28.4034495"
time_spent_string="#section-directions-trip-0 > div > div:nth-child(1) > div.xB1mrd-T3iPGc-iSfDt-n5AaSd > div.xB1mrd-T3iPGc-iSfDt-duration.delay-heavy.gm2-subtitle-alt-1 > span:nth-child(1)"
route_string="#section-directions-trip-title-0 > span"

SampleNumber=5

trafficSheet.write(0,0,"Route")
trafficSheet.write(0,1,"Time spent")
trafficSheet.write(0,2,"Sampling date")
trafficSheet.write(0,3,"Sampling time")
i=1
driver=webdriver.Chrome()
driver.get(URL)

while(i<=SampleNumber):
    try:
        driver.refresh()
        time.sleep(5)
        time_spent=driver.find_element(By.CSS_SELECTOR,time_spent_string)
        route=driver.find_element(By.CSS_SELECTOR,route_string)
        now = datetime.now()
        time.sleep(5)
        time_spent_num=time_spent.text.split()[0]
        print("Date sampling number: "+str(i))
        print("Recommended route: "+route.text)
        print("Time spent: "+time_spent.text)
        dt_string = now.strftime("%m/%d/%Y %H:%M:%S")
        print("Sampling date and time: ", dt_string)
        trafficSheet.write(i,0,route.text)
        trafficSheet.write(i,1,int(time_spent_num))
        current_date=dt_string.split()[1]
        current_time=dt_string.split()[0]
        trafficSheet.write(i,2,current_time)
        trafficSheet.write(i,3,current_date)

    except NoSuchElementException:
        print("Element not found")
        continue
    print("***")
    i+=1
trafficWorkbook.close()


from pygame import image, transform
from src.base import SCREEN_HEIGHT

GRAVITY = 0.4

class Bird:
    def __init__(self) -> None:
        self.x = 100
        self.y = SCREEN_HEIGHT / 2 - 100
        self.speed = 0
        self.downflap_image = image.load(
            "./images/redbird-downflap.png"
        ).convert_alpha()
        self.midflap_image = image.load("./images/redbird-midflap.png").convert_alpha()
        self.upflap_image = image.load("./images/redbird-upflap.png").convert_alpha()
        self.images = [self.midflap_image, self.downflap_image, self.upflap_image]
        self.rect = self.midflap_image.get_rect(center=(self.x, self.y))
        self.is_falling = True
        self.index = 0

    def draw(self, screen):
        angle = -60 if self.is_falling else 30
        image = transform.rotate(self.images[int(self.index)], angle)

        screen.blit(image, self.rect)

    def update(self):
        # Update Bird Position
        self.speed += GRAVITY
        self.rect.bottom += self.speed

        # Update Sprite
        self.index += 0.1
        if self.index >= len(self.images):
            self.index = 0

        if self.speed > 6:
            self.is_falling = True
        else:
            self.is_falling = False

    def flap(self):
        self.speed = -5

    def check_collision(self, game_asset_rect):
        return self.rect.colliderect(game_asset_rect)


"""
A simple realization of MATPOWER AC state estimation using PyPower
which can be used alongside with Python based power system steady state analysis
Sources:
    1. MATPOWER: https://github.com/MATPOWER/matpower
    2. MATPOWER SE: https://github.com/MATPOWER/mx-se
    3. PYPOWER: https://github.com/rwl/PYPOWER

The original repo for this code can be found in my GitHub: https://github.com/xuwkk/steady-state-power-system

Some code is removed if not related to the this task

Author: W XU
"""

import numpy as np
from pypower.api import *
from pypower.idx_bus import *
from pypower.idx_brch import *
from pypower.idx_gen import *
import scipy
from configs.config_mea_idx import define_mea_idx_noise
from configs.config import se_config, opt
import copy
from scipy.stats.distributions import chi2
import torch
import sys
from configs.nn_setting import nn_setting
from copy import deepcopy

def cartesian_complex_mul(A_real, A_imag, B_real, B_imag):
    """
    Using torch matmul for 2-D matrix product
    C = A * B
    return: the real and imaginary parts of C
    """
    C_real = torch.matmul(A_real, B_real) - torch.matmul(A_imag, B_imag)
    C_imag = torch.matmul(A_imag, B_real) + torch.matmul(A_real, B_imag)

    return C_real, C_imag

class SE:
    def __init__(self, case, noise_sigma, idx, fpr):

        """
        case: the instances case by calling from pypower api, e.g. case = case14()
        noise_sigma = A 1D array contains the noise std of the measurement, please refer to the format in mea_idx
        tol: the tolerance on the minimum jacobian matrix norm changes before considered as convegent
        max_it: maximum iteration
        verbose: description settings on the output
        measurement: the measurement given by each measurement type
        idx: the measurement index given by each measurement type, please refer to the format in mea_idx

        measurement type (the order matters)
        1. z = [pf, pt, pg, vang, qf, qt, qg, vmag]  (in MATPOWER-SE)
        2. z = [pf, pt, pi, vang, qf, qt, qi, vmag]  (in our current settings)
        In the future version, the selection on 1 and 2 should be added
        """

        """
        Define the grid parameter
        """

        # Case
        self.case = case

        # For analysis only
        # case['branch'][:,TAP] = 0

        case_int = ext2int(case)                                                         # Covert the start-1 to start-0 in python
        self.case_int = case_int

        # Numbers
        self.no_bus = len(case['bus'])
        self.no_brh = len(case['branch'])
        self.no_gen = len(case['gen'])
        self.no_mea = 0
        for key in idx.keys():
            self.no_mea = self.no_mea + len(idx[key])

        # Determine the bus type
        self.ref_index, pv_index, pq_index = bustypes(case_int['bus'], case_int['gen'])  # reference bus (slack bus), pv bus, and pq (load bus)
        self.non_ref_index = list(pq_index) + list(pv_index)                             # non reference bus
        self.non_ref_index.sort()

        """
        Define matrices related to measurement noise
        """
        self.noise_sigma = noise_sigma                     # std
        self.R = np.diag(noise_sigma**2)                   # R
        self.R_inv = np.diag(1/self.noise_sigma**2)        # R^-1
        self.R_inv_12 = np.diag(1/self.noise_sigma)        # R^-1/2

        self.R_inv_12_ = self.R_inv_12[:self.no_brh,:self.no_brh]  # For the active power flow

        DoF = self.no_mea - 2*(self.no_bus - 1)            # Degree of Freedom
        self.bdd_threshold = chi2.ppf(1-fpr, df = DoF)     # BDD detection threshold

        # For the power flow measurement system
        DoF_pf = self.no_brh - (self.no_bus - 1)
        self.bdd_threshold_pf = chi2.ppf(1-fpr, df = DoF_pf)

        # Hidden threshold
        fpr_attacker = 0.4
        self.bdd_threshold_attacker = chi2.ppf(1-fpr_attacker, DoF)

        """
        Incidence Matrix
        """

        # Branch Incidence Matrix
        f_bus = case_int['branch'][:, 0].astype('int')        # list of "from" buses
        t_bus = case_int['branch'][:, 1].astype('int')        # list of "to" buses
        self.Cf = np.zeros((self.no_brh,self.no_bus))         # "from" bus incidence matrix
        self.Ct = np.zeros((self.no_brh,self.no_bus))         # "to" bus incidence matrix
        for i in range(self.no_brh):
            self.Cf[i,f_bus[i]] = 1
            self.Ct[i,t_bus[i]] = 1

        self.A = self.Cf-self.Ct
        self.Ar = np.delete(self.A, self.ref_index, axis = 1) # Remove the reference bus column

        # Generator Incidence Matrix
        self.Cg = np.zeros((self.no_bus,self.no_gen))
        for i in range(self.no_gen):
            self.Cg[int(case_int['gen'][i,0]),i] = 1

        # Measurement incidence Matrix
        self.idx = idx
        no_idx_all = 4*self.no_brh + 4*self.no_bus
        self.IDX = np.zeros((self.no_mea, no_idx_all))
        _cache1 = 0
        _cache2 = 0
        for key in idx.keys():
            for _idx, _value in enumerate(idx[key]):
                self.IDX[_cache1+_idx, _cache2+_value] = 1
            _cache1 = _cache1 + len(idx[key])
            if key == 'pf' or key == 'pt' or key == 'qf' or key == 'qt':
                _cache2 = _cache2 + self.no_brh
            else:
                _cache2 = _cache2 + self.no_bus

        # Calculate the admittance matrix
        self._admittance_matrix()

        # Record the original parameters
        self.r = copy.deepcopy(case['branch'][:,BR_R])
        self.x = copy.deepcopy(case['branch'][:,BR_X])
        y = 1/(case['branch'][:,BR_R]+1j*case['branch'][:,BR_X])
        self.g = np.real(y)  # Conductance
        self.b = np.imag(y)  # Suseptance
        self.bc = case['branch'][:,BR_B]

        # Tap ratio
        self.tap = np.ones((self.no_brh,))
        tap_idx = np.where(self.case['branch'][:,TAP] != 0)[0]
        self.tap[tap_idx] = self.case['branch'][tap_idx,TAP]

        # Tap phase shift
        self.theta_s = self.case['branch'][:,SHIFT]

    def update_reactance(self,x_new):
        """
        Update reactance in self.case
        """
        self.case['branch'][:,BR_X] = x_new
        self.case_int = ext2int(self.case)

        # Update the admittance matrix
        self._admittance_matrix()

    def _admittance_matrix(self):
        """
        Calculate the Admittance matrix according to the current admittance in self.case
        """
        Ybus, Yf, Yt = makeYbus(self.case_int['baseMVA'], self.case_int['bus'], self.case_int['branch'])
        self.Ybus = scipy.sparse.csr_matrix.todense(Ybus).getA()
        self.Yf = scipy.sparse.csr_matrix.todense(Yf).getA()
        self.Yt = scipy.sparse.csr_matrix.todense(Yt).getA()

        self.Gsh = self.case['bus'][:,GS]/self.case['baseMVA']
        self.Bsh = self.case['bus'][:,BS]/self.case['baseMVA']

    def run_opf(self, verbose = False, **kwargs):
        """
        Run the optimal power flow
        """

        case_opf = copy.deepcopy(self.case)
        if 'load_active' in kwargs.keys():
            # if a new load condition is given
            case_opf['bus'][:,PD] = kwargs['load_active']
            case_opf['bus'][:,QD] = kwargs['load_reactive']
        else:
            # Use the default load condition in the case file
            pass

        if verbose:
            result = runopf(case_opf)
        else:
            result = runopf(case_opf, opt)

        return result

    def construct_mea(self, result):
        """
        Given the OPF result, construct the measurement vector
        z = [pf, pt, pi, vang, qf, qt, qi, vmag] in the current setting
        """
        pf = result['branch'][:,PF]/self.case['baseMVA']
        pt = result['branch'][:,PT]/self.case['baseMVA']
        pi = (self.Cg@result['gen'][:,PG] - result['bus'][:,PD])/self.case['baseMVA']
        vang = result['bus'][:, VA]*np.pi/180             # In radian

        qf = result['branch'][:,QF]/self.case['baseMVA']
        qt = result['branch'][:,QT]/self.case['baseMVA']
        qi = (self.Cg@result['gen'][:,QG] - result['bus'][:,QD])/self.case['baseMVA']
        vmag = result['bus'][:, VM]

        z = np.concatenate([pf, pt, pi, vang, qf, qt, qi, vmag], axis = 0)

        z = self.IDX@z   # Select the measurement
        #print(z.shape)
        #print(self.R.shape)

        z_noise = z + np.random.multivariate_normal(mean = np.zeros((self.no_mea,)), cov = self.R)
        z = np.expand_dims(z, axis = 1)
        z_noise = np.expand_dims(z_noise, axis = 1)

        vang_ref = vang[self.ref_index]
        vmag_ref = vmag[self.ref_index]

        return z, z_noise, vang_ref, vmag_ref

    def construct_v_flat(self, vang_ref, vmag_ref):

        """
        Construct a flat start voltage Given the reference bus voltage
        vmag_ref: the reference bus voltage magnitude from result
        vang: the reference bus voltage phase angle from result
        """
        vang_flat = np.zeros((self.no_bus,))
        vmag_flat = np.ones((self.no_bus,))

        vang_flat[self.ref_index] = vang_ref
        vmag_flat[self.ref_index] = vmag_ref

        return vang_flat, vmag_flat

    def h_x_pypower(self, v):
        """
        Estimate the measurement from the state: z_est = h(v)
        v is complex power
        z = [pf, pt, pi, vang, qf, qt, qi, vmag] in the current setting
        """
        #print(( np.diag(self.Cf@v)).shape)
        #print((np.conj(self.Yf)).shape)
        #print((np.conj(v).shape))

        # print(np.linalg.norm(self.Yf, 2))

        sf = np.diag(self.Cf@v)@np.conj(self.Yf)@np.conj(v)    # "from" complex power flow
        st = np.diag(self.Ct@v)@np.conj(self.Yt)@np.conj(v)    # "to" complex power flow
        si = np.diag(v)@np.conj(self.Ybus)@np.conj(v)          # complex power injection
        vang = np.angle(v)
        vmag = np.abs(v)

        pf = np.real(sf)
        pt = np.real(st)
        pi = np.real(si)
        qf = np.imag(sf)
        qt = np.imag(st)
        qi = np.imag(si)

        z_est = np.concatenate([pf, pt, pi, vang, qf, qt, qi, vmag], axis = 0)
        z_est = np.expand_dims(z_est, axis = 1)
        z_est = self.IDX@z_est

        return z_est

    def jacobian(self, v_est):
        """
        Given the stationary state, Calculate the Jacobian matrix.
        Return: the Jacobian matrix of full measurement
        """
        #
        # Compute the Jacobian matrix
        [dsi_dvmag, dsi_dvang] = dSbus_dV(self.Ybus, v_est)   # si w.r.t. v
        [dsf_dvang, dsf_dvmag, dst_dvang, dst_dvmag, _, _] = dSbr_dV(self.case_int['branch'], self.Yf, self.Yt, v_est)  # sf w.r.t. v

        dpf_dvang = np.real(dsf_dvang)
        dqf_dvang = np.imag(dsf_dvang)
        dpf_dvmag = np.real(dsf_dvmag)
        dqf_dvmag = np.imag(dsf_dvmag)

        dpt_dvang = np.real(dst_dvang)
        dqt_dvang = np.imag(dst_dvang)
        dpt_dvmag = np.real(dst_dvmag)
        dqt_dvmag = np.imag(dst_dvmag)

        dpi_dvang = np.real(dsi_dvang)
        dqi_dvang = np.imag(dsi_dvang)
        dpi_dvmag = np.real(dsi_dvmag)
        dqi_dvmag = np.imag(dsi_dvmag)

        dvang_dvang = np.eye(self.no_bus)
        dvang_dvmag = np.zeros((self.no_bus, self.no_bus))

        dvmag_dvang = np.zeros((self.no_bus, self.no_bus))
        dvmag_dvmag = np.eye(self.no_bus)

        # z = [pf, pt, pi, vang, qf, qt, qi, vmag] in the current setting
        J = np.block([
            [dpf_dvang,    dpf_dvmag],
            [dpt_dvang,    dpt_dvmag],
            [dpi_dvang,    dpi_dvmag],
            [dvang_dvang,  dvang_dvmag],
            [dqf_dvang,    dqf_dvmag],
            [dqt_dvang,    dqt_dvmag],
            [dqi_dvang,    dqi_dvmag],
            [dvmag_dvang,  dvmag_dvmag]
        ])
        # # Remove the reference bus
        # J = np.block([
        #     [dpf_dvang[:,self.non_ref_index],    dpf_dvmag[:,self.non_ref_index]],
        #     [dpt_dvang[:,self.non_ref_index],    dpt_dvmag[:,self.non_ref_index]],
        #     [dpi_dvang[:,self.non_ref_index],    dpi_dvmag[:,self.non_ref_index]],
        #     [dvang_dvang[:,self.non_ref_index],  dvang_dvmag[:,self.non_ref_index]],
        #     [dqf_dvang[:,self.non_ref_index],    dqf_dvmag[:,self.non_ref_index]],
        #     [dqt_dvang[:,self.non_ref_index],    dqt_dvmag[:,self.non_ref_index]],
        #     [dqi_dvang[:,self.non_ref_index],    dqi_dvmag[:,self.non_ref_index]],
        #     [dvmag_dvang[:,self.non_ref_index],  dvmag_dvmag[:,self.non_ref_index]]
        # ])

        J = np.array(J)         # Force convert to numpy array

        return J     # This J is not full column rank as the reference bus is not removed.

    def ac_se_pypower(self, z_noise, vang_ref, vmag_ref, is_honest = True, **kwargs):
        """
        AC-SE based on pypower
        v_initial: initial gauss on the
        is_honest:
        If True, then honest state estimation is used where the Jacobian matrix is updated at each iteration
        If False, then dishonest state estimation is used where the Jacobian matrix is fixed at the initial point.
        """
        """
        Verbose
        """
        if len(kwargs.keys()) == 0:
            # Default verbose
            tol = 1e-5,
            max_it = 100
            verbose = 0

        else:
            tol = kwargs['config']['tol']
            max_it = kwargs['config']['max_it']
            verbose = kwargs['config']['verbose']

        """
        Initialization
        """
        is_converged = 0
        ite_no = 0
        vang_est, vmag_est = self.construct_v_flat(vang_ref, vmag_ref)      # Flat start state
        v_est = vmag_est*np.exp(1j*vang_est)             # (no_bus, )

        # For the first run and also the dishonest Jacobian, the below value will never change
        J = self.jacobian(v_est)       # Jacobian matrix on the flat state
        # Remove the reference columns (for both angle and magnitude)
        J = np.delete(J, [self.ref_index, self.ref_index+self.no_bus], 1)
        J = self.IDX@J          # Select the measurement
        # Update rule: x := x_0 + (Jx0^T * R^-1 * Jx0)^-1 * Jx0^T * R^-1 * (z-h(x_0))
        G = J.T@self.R_inv@J
        G_inv = np.linalg.inv(G)

        """
        Gauss-Newton Iteration
        """

        while is_converged == False and ite_no < max_it:
            # Update iteration counter
            ite_no += 1

            # Compute estimated measurement
            z_est = self.h_x_pypower(v_est)       # z is 2D array (no_mea, 1)

            if is_honest == False:
                # No need to update the Jacobian
                #print('not update')
                pass
            else:
                #print('update')
                # Update the Jacobian
                J = self.jacobian(v_est)         # It is repeated for the first run on v flat!
                J = np.delete(J, [self.ref_index, self.ref_index+self.no_bus], 1)
                J = self.IDX@J          # Select the measurement
                # Update rule: x := x_0 + (Jx0^T * R^-1 * Jx0)^-1 * Jx0^T * R^-1 * (z-h(x_0))
                G = J.T@self.R_inv@J
                G_inv = np.linalg.inv(G)

            # Test observability
            rankG = np.linalg.matrix_rank(G)
            if rankG < G.shape[0]:
                print(f'The current measurement setting is not observable.')
                break

            F = J.T@self.R_inv@(z_noise-z_est)
            dx = (G_inv@F).flatten()  # Note that the voltages are 1D array
            normF = np.linalg.norm(F, np.inf)

            if verbose == 0:
                pass
            else:
                print(f'iteration {ite_no} norm of mismatch: {np.round(normF,6)}')

            # Terminate condition
            if normF < tol:
                is_converged = True

            # Update
            vang_est[self.non_ref_index] = vang_est[self.non_ref_index] + dx[:len(self.non_ref_index)]
            vmag_est[self.non_ref_index] = vmag_est[self.non_ref_index] + dx[len(self.non_ref_index):]
            v_est = vmag_est*np.exp(1j*vang_est)

        return v_est

    def bdd_residual(self, z_noise, v_est):
        """
        Find the residual of chi^2 detector given the estimated state
        """
        # Find z_est
        z_est = self.h_x_pypower(v_est)

        return ((z_noise-z_est).T@self.R_inv@(z_noise-z_est))[0,0]

    """
    Recovery
    """
    def gen_torch_constant(self):
        """
        Convert the numpy array into torch array for recovery
        """
        self.device = nn_setting['device']
        # print(f'device: {self.device}')
        self.Cf_torch = torch.tensor(self.Cf).float().to(self.device)
        self.Ybus_real_torch = torch.real(torch.tensor(self.Ybus)).float().to(self.device)
        self.Ybus_imag_torch = -torch.imag(torch.tensor(self.Ybus)).float().to(self.device)          # negative because of the conjugate
        self.Yf_real_torch = torch.real(torch.tensor(self.Yf)).float().to(self.device)
        self.Yf_imag_torch = -torch.imag(torch.tensor(self.Yf)).float().to(self.device)              # negative because of the conjugate
        self.R_inv_torch = torch.tensor(self.R_inv).float().to(self.device)
        self.R_inv_12_torch = torch.tensor(self.R_inv_12).float().to(self.device)

    def cartesian_state_to_measurement(self, v_real, v_imag):
        '''
        Using torch and cartesian format state to calculate the measurement
        This is further used in reconstruction measurement,
        and v_real and v_imag should be decision variable

        v_real: real part of complex state
        v_imag: imaginary part of complex state

        ref: matpower manual: https://matpower.org/docs/MATPOWER-manual.pdf
        '''
        v_diag_real = torch.diag((v_real))
        v_diag_imag = torch.diag((v_imag))
        v_real = v_real.unsqueeze(1)                # extend the dimension by one
        v_imag = -v_imag.unsqueeze(1)               # negative because of the conjugate

        # power injections
        v_diag_Ybus_real, v_diag_Ybus_imag = cartesian_complex_mul(v_diag_real, v_diag_imag, self.Ybus_real_torch, self.Ybus_imag_torch)
        P_inj, Q_inj = cartesian_complex_mul(v_diag_Ybus_real,v_diag_Ybus_imag,v_real,v_imag)

        # power flows
        Cf_v_real, Cf_v_imag = torch.matmul(self.Cf_torch,v_real), torch.matmul(self.Cf_torch,-v_imag)         # the minus here is because there is no conjugate here
        Cf_v_diag_real, Cf_v_diag_imag = torch.diag(Cf_v_real.squeeze(1)), torch.diag(Cf_v_imag.squeeze(1))
        Yf_V_real, Yf_V_imag = cartesian_complex_mul(self.Yf_real_torch, self.Yf_imag_torch, v_real, v_imag)
        P_from, Q_from = cartesian_complex_mul(Cf_v_diag_real, Cf_v_diag_imag, Yf_V_real, Yf_V_imag)
        #return v_diag_Ybus_real
        return torch.cat([P_from,P_inj,Q_from,Q_inj], axis = 0)

    """
    MTD
    """

    # max rank MTD (aka, random MTD)
    def random_mtd(self, x_max_ratio, x_min_ratio = 0.05, dfacts_index = 'full'):
        # the reactance should be inside the range
        # we note that the reactance should not be too small
        # e.g. x_max_ratio = 0.2
        # e.g. x_min_ratio = 0.05
        # dfacts_index is the branches installed with D-FACTS devices : START FROM 0!!!!!
        # assume all the branches have D-FACTS devices

        if dfacts_index == 'full':
            dfacts_index = np.arange(self.no_brh)

        pertur_ratio = x_min_ratio + (x_max_ratio - x_min_ratio)*np.random.rand(self.no_brh) # [0.05,0.2] positive
        pos_neg = np.random.choice([-1,1], self.no_brh)                                      # randomly [-0.2,-0.05] or [0.05,0.2]
        pertur_ratio_ = pertur_ratio * pos_neg

        x_mtd = deepcopy(self.x)

        x_mtd = x_mtd*(1+pertur_ratio_)

        # set 0 to the non D-FACTS devices
        non_dfact_branch = list(set(np.arange(self.no_brh)) - set(dfacts_index))
        x_mtd[non_dfact_branch] = self.x[non_dfact_branch]

        assert np.all(np.abs(x_mtd[dfacts_index]-self.x[dfacts_index])/self.x[dfacts_index] >= x_min_ratio) ## test
        assert np.all(np.abs(x_mtd[dfacts_index]-self.x[dfacts_index])/self.x[dfacts_index] <= x_max_ratio)

        return x_mtd

    # Construct matrices for optimization problems

    def H_v_hat(self, v_est):
        """
        The Jacobian matrix from active power flow to voltage phase angle
        v_est is the previous state (before MTD is triggered)
        """

        # Method 1, directly calculate from pypower function dSbr_dV
        [dsf_dvang_fun, _, _, _, _, _] = dSbr_dV(self.case_int['branch'], self.Yf, self.Yt, v_est)  # sf w.r.t. v
        dpf_dvang_fun = np.delete(np.real(dsf_dvang_fun), self.ref_index, axis = 1)
        # Method 2, calculate from the explicit equation
        # vmag = np.abs(v_est)
        # vang = np.angle(v_est)
        # V =  np.diag((self.Cf@vmag) * (self.Ct@vmag))
        # G = np.diag(self.g)
        # B = np.diag(self.b)
        # Ars = np.diag(np.sin(self.A@vang))@self.Ar
        # Arc = np.diag(np.cos(self.A@vang))@self.Ar
        # C = V@G@Ars
        # dpf_dvang_equ = C-V@B@Arc

        return np.array(dpf_dvang_fun)

    def H_v_hat_robust(self, v_est):
        """
        Generate matrices for MTD effectiveness
        """

        vmag = np.abs(v_est)
        vang = np.angle(v_est)
        V =  np.diag((self.Cf@vmag) * (self.Ct@vmag))
        G = np.diag(self.g)
        Ars = np.diag(np.sin(self.A@vang))@self.Ar
        C = V@G@Ars

        Arc = np.diag(1/self.tap) @ np.diag(np.cos(self.A@vang))@self.Ar

        return V, C, Ars, Arc

    def H_v_flat(self, vang_ref, vmag_ref):
        """
        Generate matrices for MTD effectiveness
        """

        # Construct the flat voltage
        vang_flat, vmag_flat = self.construct_v_flat(vang_ref, vmag_ref)
        V =  np.diag((self.Cf@vmag_flat) * (self.Ct@vmag_flat))
        G = np.diag(self.g)
        Ars = np.diag(np.sin(self.A@vang_flat))@self.Ar
        Arc = np.diag(1/self.tap) @ np.diag(np.cos(self.A@vang_flat))@self.Ar   # Consider the influence of tap ratio
        C = V@G@Ars

        return V, C, Ars, Arc

    def mtd_matrix(self,Jr_N):

        #Jr_N: normalized reduced jacobian matrix, calculate the projection matrix
        Pr_N = Jr_N@np.linalg.inv(Jr_N.T@Jr_N)@Jr_N.T
        Sr_N = np.eye(self.no_brh) - Pr_N

        return Pr_N, Sr_N

    def hidden_matrix(self, v_est):

        """
        Matrices for MTD hiddenness.
        Calculate the S_v_hat in the objective of hiddenness, also normalized w.r.t. the measurement noise
        """

        """
        S_v_hat_N
        """
        # Calculate the Jacobian matrix
        H_v_est = self.jacobian(v_est=v_est)   # Full Jacobian matrix
        H_v_est = self.IDX@H_v_est             # Selected Jacobian matrix
        H_v_est_N = self.R_inv_12@H_v_est      # Normalization

        # Calculate residual matrix
        S_v_hat_N = np.eye(self.no_mea) - H_v_est_N@np.linalg.inv(H_v_est_N.T@H_v_est_N)@H_v_est_N.T

        """
        H_b_N
        """
        tau = copy.deepcopy(self.tap)
        theta_s = copy.deepcopy(self.theta_s)

        # Derivative of admittance matrix to susceptance (diag)
        delta_Yff = np.diag(1j*1/(tau**2))
        delta_Yft = np.diag(-1j/(tau*np.exp(-1j*theta_s)))
        delta_Ytf = np.diag(-1j*1/(tau*np.exp(1j*theta_s)))
        delta_Ytt = np.diag(1j*np.ones(self.no_brh))

        Cf = copy.deepcopy(self.Cf)
        Ct = copy.deepcopy(self.Ct)
        v_star = np.conjugate(v_est)

        # Power Injections
        dSbus_db_1 = np.diag(Cf@v_star)@np.conjugate(delta_Yff) + np.diag(Ct@v_star)@np.conjugate(delta_Yft)
        dSbus_db_2 = np.diag(Cf@v_star)@np.conjugate(delta_Ytf) + np.diag(Ct@v_star)@np.conjugate(delta_Ytt)
        dSbus_db = np.diag(v_est)@(Cf.T@dSbus_db_1+Ct.T@dSbus_db_2)

        dPbus_db = np.real(dSbus_db)
        dQbus_db = np.imag(dSbus_db)

        # Power Flows
        dSf_db = np.diag(Cf@v_est)@(np.diag(Cf@v_star)@np.conjugate(delta_Yff) + np.diag(Ct@v_star)@np.conjugate(delta_Yft))
        dPf_db = np.real(dSf_db)
        dQf_db = np.imag(dSf_db)

        # Concat:
        # NOTE: we only write the case when the measurement is composed by 'HALF_RTU': [pf,pi,qf,qi]
        H_b = np.concatenate([dPf_db,dPbus_db,dQf_db,dQbus_db], axis = 0)

        H_b_N = self.R_inv_12@H_b   # Normalization

        H_hid = S_v_hat_N@H_b_N

        return H_hid

"""
An example
"""
if __name__ == "__main__":
    case = case14()

    # Define measurement idx
    mea_idx, no_mea, noise_sigma = define_mea_idx_noise(case, 'FULL')
    # Instance the state estimation class
    se = SE(case, noise_sigma=noise_sigma, idx=mea_idx, fpr = 0.02)

    # Run OPF to get the measurement
    opt = ppoption()              # OPF options
    opt['VERBOSE'] = 0
    opt['OUT_ALL'] = 0
    opt['OPF_FLOW_LIM'] = 1       # Constraint on the active power flow
    result = runopf(case, opt)

    print(result['success'])

    # Construct the measurement
    z, z_noise, vang_ref, vmag_ref = se.construct_mea(result) # Get the measurement

    # Run AC-SE
    se_config['verbose'] = 1
    v_est, _ = se.ac_se_pypower(z_noise, vang_ref, vmag_ref, config = se_config)
    residual = se.bdd_residual(z_noise, v_est)
    print(f'BDD threshold: {se.bdd_threshold}')
    print(f'residual: {residual}')


# code is based on https://github.com/katerakelly/pytorch-maml
import torchvision
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torch
from torch.utils.data import DataLoader,Dataset
import random
import os
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data.sampler import Sampler
import cv2
#import videotransforms
from torchvideotransforms import video_transforms, volume_transforms
from scipy.ndimage.interpolation import rotate
import json, pickle
from numpy.random import randint
from multiprocessing.dummy import Pool

def imshow(img):
    npimg = img.numpy()
    plt.axis("off")
    plt.imshow(np.transpose(npimg,(1,2,0)))
    plt.show()

class Rotate(object):
    def __init__(self, angle):
        self.angle = angle
    def __call__(self, x, mode="reflect"):
        x = rotate(x, self.angle, reshape=False)
        #x = x.rotate(self.angle)
        return x

def omniglot_character_folders():
    data_folder = '../datas/omniglot_resized/'

    character_folders = [os.path.join(data_folder, family, character) \
                for family in os.listdir(data_folder) \
                if os.path.isdir(os.path.join(data_folder, family)) \
                for character in os.listdir(os.path.join(data_folder, family))]
    random.seed(1)
    random.shuffle(character_folders)

    num_train = 1200
    metatrain_character_folders = character_folders[:num_train]
    metaval_character_folders = character_folders[num_train:]

    return metatrain_character_folders,metaval_character_folders

def validate_videos_class(videos_per_class, path_frames, num_snippets, train_num, test_num, type_text, type_dataset):
        new_vid_per_class = {}
        num_vid_total = 0
        for i, (c, videos_c) in enumerate(videos_per_class.items()):
            list_vid = []
            for vid_obj in videos_c:
                if type(vid_obj) == dict:
                    vid_name = vid_obj['id']
                else:
                    vid_name = vid_obj
                path_video = os.path.join(path_frames, vid_name)
                if type_dataset == 'epicKitchens':
                    frames = vid_obj['frames']
                else:
                    frames = os.listdir(path_video)
                    frames = [f for f in frames if f.endswith(('.jpg', '.png'))]
                num_frames = len(frames)
                if num_frames >= num_snippets:
                    if type_text == 'label':
                        text = vid_obj['label']
                    else:
                        text = c
                    text = text.lower().replace('(','').replace(')','').replace('[','').replace(']','').replace('kercief','kerchief')
                    if type_dataset == 'epicKitchens':
                        list_vid.append({'video': vid_name, 'label': text, 'start_frame': vid_obj['start_frame'], 'stop_frame': vid_obj['stop_frame'], 'frames': vid_obj['frames']})
                    else:
                        list_vid.append({'video': vid_name, 'label': text})
            if len(list_vid) >= train_num + test_num:
                new_vid_per_class[c] = list_vid
                num_vid_total += len(list_vid)
        return new_vid_per_class, num_vid_total

def get_data_classes(path, path_frames, num_snippets, train_num, test_num, type_text, type_dataset):
    try:
        with open(path, 'r') as f:
            videos_per_class = json.load(f)
    except:
        with open(path, 'rb') as f:
            videos_per_class = pickle.load(f)
    videos_per_class, num_vid_total = validate_videos_class(videos_per_class, path_frames, num_snippets, train_num, test_num, type_text, type_dataset)
    return videos_per_class, num_vid_total

class S2Sv2Task(object):
    # This class is for task generation for both meta training and meta testing.
    # For meta training, we use all 20 samples without valid set (empty here).
    # For meta testing, we use 1 or 5 shot samples for training, while using the same number of samples for validation.
    # If set num_samples = 20 and chracter_folders = metatrain_character_folders, we generate tasks for meta training
    # If set num_samples = 1 or 5 and chracter_folders = metatest_chracter_folders, we generate tasks for meta testing
    def __init__(self, videos_per_class, num_classes, train_num,test_num, path_frames, num_snippets, num_episodes):

        self.path_frames = path_frames
        self.num_snippets = num_snippets
        self.num_classes = num_classes
        self.train_num = train_num
        self.test_num = test_num
        train_labels_by_episode = []
        train_roots_by_episode = []
        valid_labels_by_episode = []
        valid_roots_by_episode = []
        for i in range(num_episodes):
            list_videos_per_class = self.get_data(videos_per_class, num_classes)
            train_labels, train_roots, valid_labels, valid_roots = self.get_labels_roots(list_videos_per_class,train_num,test_num)
            train_labels_by_episode.append(train_labels)
            train_roots_by_episode.append(train_roots)
            valid_labels_by_episode.append(valid_labels)
            valid_roots_by_episode.append(valid_roots)

        self.num_vid = len(train_labels_by_episode[0]) + len(valid_labels_by_episode[0])

        self.roots = {'train': train_roots_by_episode, 'test': valid_roots_by_episode}
        self.labels = {'train': train_labels_by_episode, 'test': valid_labels_by_episode}

    # def validate_videos_class(self, videos_per_class):
    #     new_vid_per_class = {}
    #     num_vid_total = 0
    #     for i, (c, videos_c) in enumerate(videos_per_class.items()):
    #         list_vid = []
    #         for vid_obj in videos_c:
    #             vid_name = vid_obj['id']
    #             path_video = os.path.join(self.path_frames, vid_name)
    #             num_frames = len(os.listdir(path_video))
    #             if num_frames >= self.num_snippets:
    #                 list_vid.append(vid_name)
    #         if len(list_vid) >= self.train_num + self.test_num:
    #             new_vid_per_class[c] = list_vid
    #             num_vid_total += len(list_vid)
    #     return new_vid_per_class, num_vid_total

    # def get_data(self, path, num_classes):
    #     with open(path, 'r') as f:
    #         videos_per_class = json.load(f)

    #     videos_per_class, num_vid_total = self.validate_videos_class(videos_per_class)

    #     list_classes = list(videos_per_class.keys())
    #     # print('Num training classes: ', len(list_classes))
    #     new_list_classes = random.sample(list_classes,num_classes)
    #     # print("new_list_classes: ",new_list_classes)
    #     dict_keys = {c:videos_per_class[c] for c in new_list_classes}

    #     return dict_keys, num_vid_total

    def get_data(self, videos_per_class, num_classes):
        list_classes = list(videos_per_class.keys())
        # print('Num training classes: ', len(list_classes))
        new_list_classes = random.sample(list_classes,num_classes)
        # print("new_list_classes: ",new_list_classes)
        dict_keys = {c:videos_per_class[c] for c in new_list_classes}

        return dict_keys

    def get_labels_roots(self, list_videos_per_class, train_num,test_num):
        train_labels, valid_labels = [], []
        train_videos, valid_videos = [], []
        for i, (c, videos_c) in enumerate(list_videos_per_class.items()):
            videos_ran = random.sample(videos_c, len(videos_c))
            random.shuffle(videos_ran)
            train_videos.extend(videos_ran[:train_num])
            valid_videos.extend(videos_ran[train_num:train_num+test_num])

            train_label_c = i*np.ones(train_num)
            train_labels.extend(train_label_c)

            valid_label_c = i*np.ones(test_num)
            valid_labels.extend(valid_label_c)

        return train_labels, train_videos, valid_labels, valid_videos

class KineticsTask(object):
    # This class is for task generation for both meta training and meta testing.
    # For meta training, we use all 20 samples without valid set (empty here).
    # For meta testing, we use 1 or 5 shot samples for training, while using the same number of samples for validation.
    # If set num_samples = 20 and chracter_folders = metatrain_character_folders, we generate tasks for meta training
    # If set num_samples = 1 or 5 and chracter_folders = metatest_chracter_folders, we generate tasks for meta testing
    def __init__(self, videos_per_class, num_classes, train_num,test_num, path_frames, num_snippets, num_episodes):

        self.path_frames = path_frames
        self.num_snippets = num_snippets
        self.num_classes = num_classes
        self.train_num = train_num
        self.test_num = test_num
        train_labels_by_episode = []
        train_roots_by_episode = []
        valid_labels_by_episode = []
        valid_roots_by_episode = []
        for i in range(num_episodes):
            list_videos_per_class = self.get_data(videos_per_class, num_classes)
            train_labels, train_roots, valid_labels, valid_roots = self.get_labels_roots(list_videos_per_class,train_num,test_num)
            train_labels_by_episode.append(train_labels)
            train_roots_by_episode.append(train_roots)
            valid_labels_by_episode.append(valid_labels)
            valid_roots_by_episode.append(valid_roots)

        self.num_vid = len(train_labels_by_episode[0]) + len(valid_labels_by_episode[0])

        self.roots = {'train': train_roots_by_episode, 'test': valid_roots_by_episode}
        self.labels = {'train': train_labels_by_episode, 'test': valid_labels_by_episode}

    def get_data(self, videos_per_class, num_classes):
        list_classes = list(videos_per_class.keys())
        # print('Num training classes: ', len(list_classes))
        new_list_classes = random.sample(list_classes,num_classes)
        # print("new_list_classes: ",new_list_classes)
        dict_keys = {c:videos_per_class[c] for c in new_list_classes}

        return dict_keys

    def get_labels_roots(self, list_videos_per_class, train_num,test_num):
        train_labels, valid_labels = [], []
        train_videos, valid_videos = [], []
        for i, (c, videos_c) in enumerate(list_videos_per_class.items()):
            videos_ran = random.sample(videos_c, len(videos_c))
            random.shuffle(videos_ran)
            train_videos.extend(videos_ran[:train_num])
            valid_videos.extend(videos_ran[train_num:train_num+test_num])

            train_label_c = i*np.ones(train_num)
            train_labels.extend(train_label_c)

            valid_label_c = i*np.ones(test_num)
            valid_labels.extend(valid_label_c)

        return train_labels, train_videos, valid_labels, valid_videos

class VideoRecord(object):
    def __init__(self, row):
        self._data = row

    @property
    def frames(self):
        return self._data['frames']

    @property
    def num_frames(self):
        return len(self._data['frames'])

    @property
    def path_video(self):
        return self._data['path_video']

    @property
    def label(self):
        return int(self._data['label'])

class FewShotVideoDatasetV2(Dataset):

    def __init__(self, root_path, task, type_dataset, split='train',
                 num_segments=3, new_length=1, modality='RGB',
                 transform=None, random_shift=True, test_mode=False,
                 remove_missing=False, dense_sample=False, twice_sample=False, workers = None, shuffle=True):

        self.root_path = root_path
        self.type_dataset = type_dataset
        self.task = task
        self.split = split
        self.video_root = self.task.roots[split]
        self.labels = self.task.labels[split]
        self.num_segments = num_segments
        self.new_length = new_length
        self.modality = modality
        self.transform = transform
        self.random_shift = random_shift
        self.test_mode = test_mode
        self.remove_missing = remove_missing
        self.dense_sample = dense_sample  # using dense sample as I3D
        self.twice_sample = twice_sample  # twice sample for more validation
        if self.dense_sample:
            print('=> Using dense sample for the dataset...')
        if self.twice_sample:
            print('=> Using twice sample for the dataset...')

        if self.modality == 'RGBDiff':
            self.new_length += 1  # Diff needs one more image to calculate diff

        self.pool = workers
        self.shuffle = shuffle

        # self._parse_list()

    # def _parse_list(self):
    #     self.video_list = list()
    #     for i, path_vid in enumerate(self.video_root):
    #         video_path = os.path.join(self.root_path, path_vid)
    #         frames = os.listdir(video_path)
    #         label = self.labels[i]
    #         self.video_list.append(VideoRecord({'frames': frames,'label': label, 'path_video': video_path}))
    #     print('video number:%d' % (len(self.video_list)))

    def _load_image(self, frame_path, path_video):
        if self.modality == 'RGB' or self.modality == 'RGBDiff':
            try:
                return [Image.open(os.path.join(path_video, frame_path)).convert('RGB')]
            except Exception:
                print('error loading image:', os.path.join(path_video, frame_path))
                return [Image.open(os.path.join(path_video, frame_path)).convert('RGB')]
        elif self.modality == 'Flow':
            if self.image_tmpl == 'flow_{}_{:05d}.jpg':  # ucf
                x_img = Image.open(os.path.join(path_video, frame_path)).convert(
                    'L')
                y_img = Image.open(os.path.join(path_video, frame_path)).convert(
                    'L')
            elif self.image_tmpl == '{:06d}-{}_{:05d}.jpg':  # something v1 flow
                x_img = Image.open(os.path.join(path_video, '{:06d}'.format(int(frame_path)))).convert('L')
                y_img = Image.open(os.path.join(path_video, '{:06d}'.format(int(frame_path)))).convert('L')
            else:
                try:
                    # idx_skip = 1 + (idx-1)*5
                    flow = Image.open(os.path.join(self.root_path, frame_path)).convert(
                        'RGB')
                except Exception:
                    print('error loading flow file:',
                          os.path.join(self.root_path, frame_path))
                    flow = Image.open(os.path.join(self.root_path, frame_path)).convert('RGB')
                # the input flow file is RGB image with (flow_x, flow_y, blank) for each channel
                flow_x, flow_y, _ = flow.split()
                x_img = flow_x.convert('L')
                y_img = flow_y.convert('L')

            return [x_img, y_img]

    def _sample_indices(self, record):
        """

        :param record: VideoRecord
        :return: list
        """
        if self.dense_sample:  # i3d dense sample
            sample_pos = max(1, 1 + record.num_frames - 64)
            t_stride = 64 // self.num_segments
            start_idx = 0 if sample_pos == 1 else np.random.randint(0, sample_pos - 1)
            offsets = [(idx * t_stride + start_idx) % record.num_frames for idx in range(self.num_segments)]
            return np.array(offsets) + 1
        else:  # normal sample
            average_duration = (record.num_frames - self.new_length + 1) // self.num_segments
            if average_duration > 0:
                offsets = np.multiply(list(range(self.num_segments)), average_duration) + randint(average_duration,
                                                                                size=self.num_segments)
            elif record.num_frames > self.num_segments:
                offsets = np.sort(randint(record.num_frames - self.new_length + 1, size=self.num_segments))
            else:
                offsets = np.zeros((self.num_segments,))
            return offsets + 1

    def _get_val_indices(self, record):
        if self.dense_sample:  # i3d dense sample
            sample_pos = max(1, 1 + record.num_frames - 64)
            t_stride = 64 // self.num_segments
            start_idx = 0 if sample_pos == 1 else np.random.randint(0, sample_pos - 1)
            offsets = [(idx * t_stride + start_idx) % record.num_frames for idx in range(self.num_segments)]
            return np.array(offsets) + 1
        else:
            if record.num_frames > self.num_segments + self.new_length - 1:
                tick = (record.num_frames - self.new_length + 1) / float(self.num_segments)
                offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.num_segments)])
            else:
                offsets = np.zeros((self.num_segments,))
            return offsets + 1

    def _get_test_indices(self, record):
        if self.dense_sample:
            sample_pos = max(1, 1 + record.num_frames - 64)
            t_stride = 64 // self.num_segments
            start_list = np.linspace(0, sample_pos - 1, num=10, dtype=int)
            offsets = []
            for start_idx in start_list.tolist():
                offsets += [(idx * t_stride + start_idx) % record.num_frames for idx in range(self.num_segments)]
            return np.array(offsets) + 1
        elif self.twice_sample:
            tick = (record.num_frames - self.new_length + 1) / float(self.num_segments)

            offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.num_segments)] +
                               [int(tick * x) for x in range(self.num_segments)])

            return offsets + 1
        else:
            tick = (record.num_frames - self.new_length + 1) / float(self.num_segments)
            offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.num_segments)])
            return offsets + 1

    def get_elem(self, index):
        vid_obj = self.episode_path_vid[index]
        path_vid = vid_obj['video']
        label_text = vid_obj['label']
        label = self.episode_label[index]

        video_path = os.path.join(self.root_path, path_vid)

        if self.type_dataset == 'epicKitchens':
            frames = vid_obj['frames']
            frames_num = [i.split('.')[0].split('_')[1] for i in frames]
        else:
            frames = os.listdir(video_path)
            frames = [f for f in frames if f.endswith(('.jpg', '.png'))]
            frames_num = [i.split('.')[0] for i in frames]
        idxs = sorted(range(len(frames_num)), key=lambda k: float(frames_num[k]))
        frames = [frames[i] for i in idxs]

        record = VideoRecord({'frames': frames,'label': label, 'path_video': video_path})

        if not self.test_mode:
            segment_indices = self._sample_indices(record) if self.random_shift else self._get_val_indices(record)
        else:
            segment_indices = self._get_test_indices(record)

        img, label = self.get(record, segment_indices)
        return img, label, label_text

    def __getitem__(self, index):
        self.episode_path_vid = self.video_root[index]
        self.episode_label = self.labels[index]

        list_img_episode = self.pool.map(self.get_elem, list(range(len(self.episode_label))))
        if self.shuffle:
            random.shuffle(list_img_episode)

        imgs_episode, labels_episode, labels_text_episode = zip(*list_img_episode)
        imgs_episode = torch.stack(imgs_episode, dim = 0)
        labels_episode = torch.tensor(labels_episode)
        return imgs_episode, labels_episode, list(labels_text_episode)

    def get(self, record, indices):

        images = list()
        frames = record.frames
        for seg_ind in indices:
            p = int(seg_ind)
            for i in range(self.new_length):
                seg_imgs = self._load_image(frames[p-1],record.path_video)
                images.extend(seg_imgs)

                if p < record.num_frames:
                    p += 1

        process_data = self.transform(images)
        return process_data, record.label

    def __len__(self):
        return len(self.video_root)

class FewShotVideoDataset(Dataset):

    def __init__(self, path_frames, task, num_snippets = 20, split='train', transform=None, rotate=None, target_transform=None):
        self.transform = transform
        self.target_transform = target_transform
        self.task = task
        self.split = split
        self.video_root = self.task.roots[split]
        self.labels = self.task.labels[split]
        self.num_snippets = num_snippets
        self.rotate = rotate
        self.path_frames = path_frames

    def __getitem__(self, idx):

        video_path = os.path.join(self.path_frames, self.video_root[idx])
        frames = os.listdir(video_path)
        lon_snippets = int(len(frames)/self.num_snippets)
        #values = np.linspace(0, len(frames) - 1, num=self.num_snippets)
        i = 0
        while(i<self.num_snippets):
            value = random.randint(lon_snippets*i , lon_snippets*(i+1) - 1)
        #for i,value in enumerate(values):
            frame = frames[int(value)]
            frame_path = os.path.join(video_path, frame)
            image = cv2.imread(frame_path)[:, :, [2, 1, 0]]
            w,h,c = image.shape
            if w < 226 or h < 226:
                d = 226.-min(w,h)
                sc = 1+d/min(w,h)
                image = cv2.resize(image,dsize=(0,0),fx=sc,fy=sc)
            image = (image/255.)*2 - 1
            if i == 0:
                video_frames = np.zeros((self.num_snippets,) + image.shape, np.float32)
            if self.rotate is not None:
                print('Rotate video')
                image = self.rotate(image)
            video_frames[i, :, :, :] = image
            i += 1

        #image = image.resize((28,28), resample=Image.LANCZOS) # per Chelsea's implementation
        if self.transform is not None:
            video_frames = self.transform(video_frames)
        label = self.labels[idx]
        if self.target_transform is not None:
            label = self.target_transform(label)
        #video_frames = torch.from_numpy(video_frames.transpose([3,0,1,2]))
        return video_frames, torch.from_numpy(np.array([label]))

    def __len__(self):
        return len(self.video_root)

class FewShotDataset(Dataset):

    def __init__(self, task, split='train', transform=None, target_transform=None):
        self.transform = transform # Torch operations on the input image
        self.target_transform = target_transform
        self.task = task
        self.split = split
        self.image_roots = self.task.train_roots if self.split == 'train' else self.task.test_roots
        self.labels = self.task.train_labels if self.split == 'train' else self.task.test_labels

    def __len__(self):
        return len(self.image_roots)

    def __getitem__(self, idx):
        raise NotImplementedError("This is an abstract class. Subclass this class for your particular dataset.")

class Omniglot(FewShotDataset):

    def __init__(self, *args, **kwargs):
        super(Omniglot, self).__init__(*args, **kwargs)

    def __getitem__(self, idx):
        image_root = self.image_roots[idx]
        image = Image.open(image_root)
        image = image.convert('L')
        image = image.resize((28,28), resample=Image.LANCZOS) # per Chelsea's implementation
        #image = np.array(image, dtype=np.float32)
        if self.transform is not None:
            image = self.transform(image)
        label = self.labels[idx]
        if self.target_transform is not None:
            label = self.target_transform(label)
        return image, label

class ClassBalancedSampler(Sampler):
    ''' Samples 'num_inst' examples each from 'num_cl' pools
        of examples of size 'num_per_class' '''

    def __init__(self, num_per_class, num_cl, num_inst,shuffle=True):
        self.num_per_class = num_per_class
        self.num_cl = num_cl
        self.num_inst = num_inst
        self.shuffle = shuffle

    def __iter__(self):
        # return a single list of indices, assuming that items will be grouped by class
        if self.shuffle:
            batch = [[i+j*self.num_inst for i in torch.randperm(self.num_inst)[:self.num_per_class]] for j in range(self.num_cl)]
        else:
            batch = [[i+j*self.num_inst for i in range(self.num_inst)[:self.num_per_class]] for j in range(self.num_cl)]
        batch = [item for sublist in batch for item in sublist]

        if self.shuffle:
            random.shuffle(batch)
        return iter(batch)

    def __len__(self):
        return 1

def get_data_loader(task, type_dataset, num_per_class=1, split='train',shuffle=True,rotation=0, transforms = None, test_mode = False, workers = None):
    # NOTE: batch size here is # instances PER CLASS
    #normalize = transforms.Normalize(mean=[0.92206, 0.92206, 0.92206], std=[0.08426, 0.08426, 0.08426])
    #normalize = transforms.Normalize(mean=[0.92206,], std=[0.08426,])
    # transforms_video_train = transforms.Compose([
    #             videotransforms.RandomCrop(100)])

    # video_transform_list_train = [video_transforms.Resize((256, 256)),
	# 		video_transforms.RandomCrop((224, 224)),
    #         # video_transforms.Resize((100, 100)),
    #         video_transforms.RandomHorizontalFlip(),
	# 		volume_transforms.ClipToTensor()]

    # video_transform_list_test = [video_transforms.Resize((256, 256)),
	# 		video_transforms.CenterCrop((224, 224)),
    #         # video_transforms.Resize((100, 100)),
	# 		volume_transforms.ClipToTensor()]

    if split == 'train':

        #transforms_train = video_transforms.Compose(video_transform_list_train)
        transforms_train = torchvision.transforms.Compose(transforms)

        dataset = FewShotVideoDatasetV2(task.path_frames, task, type_dataset = type_dataset, split=split, num_segments=task.num_snippets,
                   new_length=1,
                   modality='RGB',
                   test_mode=test_mode,
                   transform=transforms_train, dense_sample=False, workers = workers, shuffle=shuffle)

        # dataset = FewShotVideoDataset(task.path_frames,task,num_snippets = task.num_snippets, split=split,transform = transforms_train)
        # sampler = ClassBalancedSampler(num_per_class, task.num_classes, task.train_num,shuffle=shuffle)
    else:
        # transforms_test = video_transforms.Compose(video_transform_list_test)
        transforms_test = torchvision.transforms.Compose(transforms)
        # dataset = FewShotVideoDataset(task.path_frames,task,num_snippets = task.num_snippets, split=split,transform = transforms_test)

        dataset = FewShotVideoDatasetV2(task.path_frames, task, type_dataset = type_dataset, split=split, num_segments=task.num_snippets,
                    new_length=1,
                    modality = "RGB",
                    test_mode=test_mode,
                    transform=transforms_test, dense_sample=False, workers = workers, shuffle=shuffle)

        # sampler = ClassBalancedSampler(num_per_class, task.num_classes, task.test_num,shuffle=shuffle)

    loader = DataLoader(dataset, batch_size=1, pin_memory=True)

    return loader


#!/usr/bin/env python

import numpy as np
from spatialmath import SE3, base
import math
from typing import Union
from roboticstoolbox.fknm import Angle_Axis

ArrayLike = Union[list, np.ndarray, tuple, set]

def angle_axis(T, Td):

    try:
        e = Angle_Axis(T, Td)
    except BaseException:
        e = angle_axis_python(T, Td)

    return e

def angle_axis_python(T, Td):
    e = np.empty(6)
    e[:3] = Td[:3, -1] - T[:3, -1]
    R = Td[:3, :3] @ T[:3, :3].T
    li = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])

    if base.iszerovec(li):
        # diagonal matrix case
        if np.trace(R) > 0:
            # (1,1,1) case
            a = np.zeros((3,))
        else:
            a = np.pi / 2 * (np.diag(R) + 1)
    else:
        # non-diagonal matrix case
        ln = base.norm(li)
        a = math.atan2(ln, np.trace(R) - 1) * li / ln

    e[3:] = a

    return e

def p_servo(
    wTe, wTep, gain: Union[float, ArrayLike] = 1.0, threshold=0.1, method="rpy"
):
    """
    Position-based servoing.

    Returns the end-effector velocity which will cause the robot to approach
    the desired pose.

    :param wTe: The current pose of the end-effecor in the base frame.
    :type wTe: SE3 or ndarray
    :param wTep: The desired pose of the end-effecor in the base frame.
    :type wTep: SE3 or ndarray
    :param gain: The gain for the controller. Can be vector corresponding to each
        axis, or scalar corresponding to all axes.
    :type gain: float, or array-like
    :param threshold: The threshold or tolerance of the final error between
        the robot's pose and desired pose
    :type threshold: float
    :param method: The method used to calculate the error. Default is 'rpy' -
        error in the end-effector frame. 'angle-axis' - error in the base frame
        using angle-axis method.
    :type method: string: 'rpy' or 'angle-axis'

    :returns v: The velocity of the end-effecotr which will casue the robot
        to approach wTep
    :rtype v: ndarray(6)
    :returns arrived: True if the robot is within the threshold of the final
        pose
    :rtype arrived: bool

    """

    if isinstance(wTe, SE3):
        wTe = wTe.A

    if isinstance(wTep, SE3):
        wTep = wTep.A

    if method == "rpy":
        # Pose difference
        eTep = np.linalg.inv(wTe) @ wTep
        e = np.empty(6)

        # Translational error
        e[:3] = eTep[:3, -1]

        # Angular error
        e[3:] = base.tr2rpy(eTep, unit="rad", order="zyx", check=False)
    else:
        e = angle_axis(wTe, wTep)

    if base.isscalar(gain):
        k = gain * np.eye(6)
    else:
        k = np.diag(gain)

    v = k @ e
    arrived = True if np.sum(np.abs(e)) < threshold else False

    return v, arrived


import os
import pyrender
import trimesh
import numpy as np
import imageio
import json
import cv2
import matplotlib.pyplot as plt
from scipy import io
from pytransform3d.rotations import matrix_from_axis_angle
from struct import iter_unpack
from pyquaternion import Quaternion
import shutil

def axis_angle_to_R(pose):
    T = np.zeros((4,4))
    aa = np.array(pose[3:])

    if np.nonzero(aa)[0].size > 0:
        R = matrix_from_axis_angle(aa)
        T[:3,:3] = R

    T[0,3] = pose[0]
    T[1,3] = pose[1]
    T[2,3] = pose[2]
    T[3,3] = 1

    return T

def T_to_aa(T_mat):
    aa = np.zeros(7)
    aa[:3] = T_mat[:3,3]
    try:
        q = Quaternion(matrix = T_mat[:3,:3])
    except ValueError:
        U, s, VT = np.linalg.svd(T_mat[:3,:3])
        q = Quaternion(matrix = np.dot(U, VT))
    aa[3:6] = q.axis
    aa[6] = q.angle

    return aa

def nvdu_to_ycbv(nvdu_path, ycbv_models_path, obj_id):
    # Object id to object name
    classes_lookup =\
    {
        '2': '003_cracker_box',
        '3': '004_sugar_box',
        '4': '005_tomato_soup_can',
        '5': '006_mustard_bottle',
        '8': '009_gelatin_box',
        '9': '010_potted_meat_can'
     }

    # Paths
    nvdu_obj_settings_path = nvdu_path + '/nvdu/config/object_settings/_ycb_original.json'
    nvdu_model_path = os.path.join(nvdu_path, 'nvdu/data/ycb/original', classes_lookup[str(obj_id)], 'google_16k/textured.obj')
    ycbv_model_path = os.path.join(ycbv_models_path, classes_lookup[str(obj_id)], 'textured_simple.obj')

    # Meshes
    nvdu_original_mesh = trimesh.load(nvdu_model_path)
    ycbv_mesh = trimesh.load(ycbv_model_path)

    # Relative transformation between NVDU aligned and NVDU original
    transform = json.load(open(nvdu_obj_settings_path))
    fixed_model_transform = np.asarray(transform['exported_objects'][int(obj_id)-1]['fixed_model_transform'], dtype='float32')
    aligned_to_original = np.zeros((4,4))

    # Divide by 100.0 to restore standard units
    for i, row in enumerate(fixed_model_transform):
        aligned_to_original[i,:] =  np.array(row) / 100.

    # Restore the scaler of the transformation to 1
    aligned_to_original[3, 3] = 1.

    # Invert direction of transformation
    aligned_to_original = aligned_to_original.T

    # Relative transformation between NVDU original and YCB-V
    original_to_ycbv = np.zeros((4,4))
    original_to_ycbv[0:3, 0:3] = np.eye(3)
    original_to_ycbv[0:3, 3] = nvdu_original_mesh.centroid - ycbv_mesh.centroid
    original_to_ycbv[3, 3] = 1.

    # Relative transformation between NVDU aligned and YCB-V
    aligned_to_ycbv = np.dot(aligned_to_original, original_to_ycbv)

    return aligned_to_ycbv

def get_K_tracknet(camera_K_txt_path):
    cam_K_file = open(camera_K_txt_path)
    cam_K = np.zeros((3,3))
    i = 0
    for row in cam_K_file:
        cam_K[i] = [float(val) for val in row.split(' ') if val != '\n']
        i += 1

    return cam_K

def get_K_synthetic(path):
    cam_file = json.load(open(path))
    cam_K = np.zeros((3,3))

    cam_K[0,0] = float(cam_file['fx'])
    cam_K[1,1] = float(cam_file['fy'])
    cam_K[0,2] = float(cam_file['cx'])
    cam_K[1,2] = float(cam_file['cy'])
    cam_K[2,2] = 1

    return cam_K

def get_obj_pose_tracknet(pose_path):
    pose_file = open(pose_path)
    pose = np.zeros((4,4))
    i = 0
    for row in pose_file:
        pose[i] = [float(val) for val in row.split(' ') if val != '\n']
        i += 1

    return pose

def get_obj_pose_and_K_ycbv(anns_path, obj_id):
    meta_dict = io.loadmat(anns_path)

    pose = np.zeros((4,4))
    ind_mapped = np.where(meta_dict['cls_indexes'] == int(obj_id))[0][0]
    pose[:3] = meta_dict['poses'][:,:,ind_mapped]
    pose[3,3] = 1

    return pose, meta_dict['intrinsic_matrix']

def get_obj_poses_synthetic(pose_path):
    pose_file = open(pose_path)
    poses = dict()
    i = 0

    for row in pose_file:
        pose_axis_angle = [float(val.split('\n')[0]) for val in row.split(' ') if val.split('\n')[0] != '']
        # Take only last 7 values, skip the first values related to velocities
        T = axis_angle_to_R(pose_axis_angle[-7:])
        poses[i] = T
        i += 1

    return poses

def crop_and_resize(img, data_type='uint8'):
    if len(img.shape) == 3:
        img_crop = np.zeros((720,960,3))
    else:
        img_crop = np.zeros((720,960))
    img_crop = img[:,160:1120]
    img_res = cv2.resize(img_crop, (640, 480), interpolation=cv2.INTER_NEAREST)

    return np.array(img_res, dtype=data_type)

def pad_and_resize(img, data_type='uint8', depth_factor=0):
    if len(img.shape) == 3:
        img_pad = np.zeros((960,1280,3))
    else:
        img_pad = np.zeros((960,1280))

    if depth_factor != 0:
        img_pad[:120] = 10 * depth_factor
        img_pad[840:] = 10 * depth_factor

    img_pad[120:840] = img
    img_res = cv2.resize(img_pad, (640, 480), interpolation=cv2.INTER_NEAREST)

    return np.array(img_res, dtype=data_type)

def convert_depth(path, in_format='synth', factor=1000):
    with open(path, 'rb') as dat:
        # This will give data as a 1D array
        data = list(iter_unpack('f', dat.read()))
    if in_format == 'synth' or in_format == 'real':
        depth = (np.array(data[4:])*factor).reshape(720, 1280)
        depth = pad_and_resize(depth, 'uint16', factor)
    elif in_format == 'ho3d':
        depth = (np.array(data[4:])*factor).reshape(480, 640)
        depth = np.array(np.where(depth==0, 10*factor, depth), dtype='uint16')

    return depth

def process_image(synth_img_path, tracknet_img_path, in_format):
    if in_format == 'synth' or in_format == 'real':
        synth_img = cv2.imread(synth_img_path)
        resized_img = pad_and_resize(synth_img)
        cv2.imwrite(tracknet_img_path, resized_img)
    elif in_format == 'ho3d':
        shutil.copy(synth_img_path, tracknet_img_path)

def convert_mask(path, obj_id, in_format='synth'):
    mask = imageio.imread(path)
    if in_format == 'synth' or in_format == 'real':
        mask = pad_and_resize(mask)
    converted = np.where(mask == 255, int(obj_id), 0)

    where = np.array(np.where(converted))
    x1, y1 = np.amin(where, axis=1) # top-left, row-column
    x2, y2 = np.amax(where, axis=1) # bottom-right, row-column
    bbox = f' {y1} {x1} {y2} {x2}\n' # column-row, top-left, bottom-right

    return converted, bbox

def render_pose(model_path, cam_K, obj_pose):
    os.environ['PYOPENGL_PLATFORM'] = 'egl'

    trimesh_model = trimesh.load(model_path)
    mesh = pyrender.Mesh.from_trimesh(trimesh_model)
    scene = pyrender.Scene()

    fx, fy, cx, cy = cam_K[0,0], cam_K[1,1], cam_K[0,2], cam_K[1,2]
    camera = pyrender.IntrinsicsCamera(fx, fy, cx, cy)
    cam_pose = np.identity(4)
    cam_pose[1,1] = 0
    cam_pose[2,2] = 0
    cam_pose[1,2] = 1
    cam_pose[2,1] = 1

    scene.add(mesh, pose=obj_pose)
    scene.add(camera, pose=cam_pose)

    light = pyrender.SpotLight(color=np.ones(3), intensity=10.0,
                                innerConeAngle=np.pi/16.0,
                                outerConeAngle=np.pi/6.0)
    scene.add(light, pose=cam_pose)

    r = pyrender.OffscreenRenderer(viewport_width=640, viewport_height=480,
                                   point_size=1.0)
    color, depth = r.render(scene)
    r.delete()

    return color, depth

def vis_rendering(color_img, depth_img):
    plt.figure()

    plt.subplot(1,2,1)
    plt.axis('off')
    plt.imshow(color_img)

    plt.subplot(1,2,2)
    plt.axis('off')
    plt.imshow(depth_img, cmap=plt.cm.gray_r)

    plt.show()

def vis_overlap_render(orig_img, color, ren_depth):
    img_overlapped = np.zeros_like(orig_img)
    grayscale = rgb2gray(orig_img)
    for i in range(3):
        img_overlapped[:,:,i] = np.where(ren_depth != 0, color[:,:,i], grayscale)

    plt.figure()
    plt.imshow(img_overlapped)

def vis_overlap_mask(orig_img, ren_depth):
    img_overlapped = np.zeros_like(orig_img)
    for i in range(3):
        img_overlapped[:,:,i] = np.where(ren_depth != 0, orig_img[:,:,i], 0)

    plt.figure()
    plt.imshow(img_overlapped)

def zero_pad(inp, out_dim):
    zeros = '0' * out_dim

    return (zeros + str(inp))[-out_dim:]

def matrix_to_txt(matrix, out_path, velNaN=False):
    file = open(out_path, 'w')
    for row in matrix:
        if velNaN:
            file.write('NaN ' * 6)
        for val in row:
            file.write(str(val))
            file.write(' ')
        file.write('\n')
    file.close()

def rgb2gray(rgb):

    r, g, b = rgb[:,:,0], rgb[:,:,1], rgb[:,:,2]
    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b

    return gray


from typing import List, Dict
from .base import Base

class Source(Base):
    def __init__(self, vim):
        super().__init__(vim)

        self.name = 'codeAction'
        self.kind = 'command'

    def gather_candidates(self, context: Dict) -> List[Dict]:
        result = self.vim.funcs.LanguageClient_runSync(
            'LanguageClient_textDocument_codeAction', {}) or []
        return [convert_to_candidate(item) for item in result]

def convert_to_candidate(cmd: Dict) -> Dict:
    cmd_str = '{}: {}'.format(cmd['command'], cmd['title'])
    return {
        'word': cmd_str,
        'action__command': 'call '
        'LanguageClient_FZFSinkCommand("{}")'.format(cmd_str),
    }


import csv
import shutil
from pathlib import Path
from typing import Optional

from appdirs import user_config_dir
from pkg_resources import resource_filename

from . import __about__

DEFAULT_FILE = resource_filename(__about__.__name__, "data/default_config.csv")
FILE         = (Path(user_config_dir(__about__.__project_name__)) /
                f"{__about__.__pkg_name__}.csv")
FIELDS = ["tag_search", "filter", "labels", "booru"]
SUBS   = []

def reload(path: Optional[str] = None) -> None:
    path = Path(path or FILE).expanduser()

    if not path.is_file():
        path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(DEFAULT_FILE, path)

    with open(path, "r") as file:
        SUBS.extend(list(
            csv.DictReader(
                (row for row in file if row and not row.startswith("#")),
                fieldnames = FIELDS
            )
        ))