import argparse
import json
import logging
import os
import pathlib
import pickle
from typing import Dict, Tuple, List, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.optim import Adam

from .config import (
    UFDAdaptorGlobalConfig,
    UFDAdaptorDomainConfig,
    UFDCombineFeaturesMapConfig,
    UFDClassifierConfig,
    UFDEmbeddingConfig,
)
from .data_class import UFDArguments
from .tokenization import UFDTokenizer
from .modeling import (
    UFDAdaptorGlobalModel,
    UFDAdaptorDomainModel,
    UFDCombineFeaturesMapModel,
    UFDClassifierModel,
    UFDEmbeddingModel,
)

logging.basicConfig(level=logging.DEBUG)

def create_unsupervised_models(
    cfg: UFDArguments,
) -> Tuple[UFDAdaptorDomainModel, UFDAdaptorGlobalModel, UFDCombineFeaturesMapModel]:
    """Helper function to create the unsupervised model group.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file

    Returns:
        Tuple[UFDAdaptorDomainModel, UFDAdaptorGlobalModel, UFDCombineFeaturesMapModel]:
            return the initialize unsupervised model group.
    """
    device = torch.device(cfg.device)
    adaptor_domain_config = UFDAdaptorDomainConfig(
        in_dim=cfg.train_args["in_dim"],
        dim_hidden=cfg.train_args["dim_hidden"],
        out_dim=cfg.train_args["out_dim"],
        initrange=cfg.train_args["initrange"],
    )
    adaptor_global_config = UFDAdaptorGlobalConfig(
        in_dim=cfg.train_args["in_dim"],
        dim_hidden=cfg.train_args["dim_hidden"],
        out_dim=cfg.train_args["out_dim"],
        initrange=cfg.train_args["initrange"],
    )
    combine_features_map_config = UFDCombineFeaturesMapConfig(
        embed_dim=cfg.train_args["in_dim"], initrange=cfg.train_args["initrange"]
    )
    return (
        UFDAdaptorDomainModel(adaptor_domain_config).to(device),
        UFDAdaptorGlobalModel(adaptor_global_config).to(device),
        UFDCombineFeaturesMapModel(combine_features_map_config).to(device),
    )

def load_trained_models(
    cfg: UFDArguments,
    source_domain: str,
    target_language: str,
    target_domain: str,
) -> Tuple[
    UFDAdaptorDomainModel,
    UFDAdaptorGlobalModel,
    UFDCombineFeaturesMapModel,
    UFDClassifierModel,
]:
    """Helper function to load pretrained config and model weights for both supervised and unsupervised
        models.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file
        source_domain (str): name of source domain
        target_language (str): name of target language
        target_domain (str): name of target domain

    Returns:
        Tuple[ UFDAdaptorDomainModel, UFDAdaptorGlobalModel, UFDCombineFeaturesMapModel, UFDClassifierModel, ]:
            return all supervised, unsupervised models with pretrained weights loaded.
    """
    device = torch.device(cfg.device)
    full_combi = (
        cfg.model_folder
        + "/"
        + source_domain
        + "_"
        + target_language
        + "_"
        + target_domain
    )
    adaptor_domain_model_dir = full_combi + "_adaptor_domain/"
    adaptor_global_model_dir = full_combi + "_adaptor_global/"
    maper_model_dir = full_combi + "_maper/"
    classifier_model_dir = full_combi + "_classifier/"

    config_filename = cfg.eval_args["config_filename"]
    model_filename = cfg.eval_args["model_filename"]

    adaptor_domain_config = UFDAdaptorDomainConfig.from_pretrained(
        adaptor_domain_model_dir + config_filename
    )
    adaptor_global_config = UFDAdaptorGlobalConfig.from_pretrained(
        adaptor_global_model_dir + config_filename
    )
    maper_config = UFDCombineFeaturesMapConfig.from_pretrained(
        maper_model_dir + config_filename
    )
    classifier_config = UFDClassifierConfig.from_pretrained(
        classifier_model_dir + config_filename
    )
    return (
        UFDAdaptorDomainModel.from_pretrained(
            adaptor_domain_model_dir + model_filename, config=adaptor_domain_config
        ).to(device),
        UFDAdaptorGlobalModel.from_pretrained(
            adaptor_global_model_dir + model_filename, config=adaptor_global_config
        ).to(device),
        UFDCombineFeaturesMapModel.from_pretrained(
            maper_model_dir + model_filename, config=maper_config
        ).to(device),
        UFDClassifierModel.from_pretrained(
            classifier_model_dir + model_filename, config=classifier_config
        ).to(device),
    )

def load_unlabelled(filename: str) -> List:
    """Helper function to load unlabelled dataset for unsupervised training.

    Args:
        filename (str): filename of dataset

    Returns:
        List: list of dataset by line.
    """
    data = []
    with open(filename, "r") as F:
        for line in F:
            # 0 is used in the research code
            data.append([0, line.strip()])
    return data

def load_labelled(filename: str) -> List:
    """Helper function to load labelled dataset for supervised training.

    Args:
        filename (str): filename of dataset

    Returns:
        List: list of dataset by line.
    """
    data = []
    with open(filename, "r") as F:
        for line in F:
            data.append(line.split("\t"))
    return data

def extract_embeddings(
    cfg: UFDArguments, dataset: List, tokenizer: UFDTokenizer, model: UFDEmbeddingModel
) -> List:
    """Helper function to extract embeddings with the UFD embedding model.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file
        dataset (List): list of dataset by line
        tokenizer (UFDTokenizer): UFD tokenizer class instance
        model (UFDEmbeddingModel): UFD embedding model class instance

    Returns:
        List: return list of generated embeddings.
    """
    device = torch.device(cfg.device)
    model.eval()
    embeded_data = []
    with torch.no_grad():
        for pair in dataset:
            tokens = tokenizer(pair[1])["input_ids"].to(device)
            last_layer_features = model(tokens)[0]  # 1 * length * embedding_size
            mean_features = torch.mean(last_layer_features, dim=1)  # 1 * embedding_size
            tem_label = torch.tensor([[float(pair[0])]]).to(device)
            new_pair = torch.cat((tem_label, mean_features), dim=1)
            embeded_data.append(new_pair)
    return embeded_data

def create_train_embeddings(
    cfg: UFDArguments, tokenizer: UFDTokenizer, model: UFDEmbeddingModel
) -> Dict:
    """Helper function to generate training dataset for supervised and unsupervised training.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file
        tokenizer (UFDTokenizer): UFD tokenizer class instance
        model (UFDEmbeddingModel): UFD embedding model class instance

    Returns:
        Dict: dictionary of dataset embeddings for supervised and unsupervised dataset
    """
    embeddings_dict = {}
    source_domains_list = cfg.train_args["source_domains"] + ["raw"]
    for source_domain in source_domains_list:
        if source_domain == "raw":
            # assume will only have one source language
            filepath = f"{cfg.data_folder}/{cfg.train_args['unsupervised_dataset_filename']}"
            dataset = load_unlabelled(filepath)
        else:
            filepath = \
                f"{cfg.data_folder}/{cfg.train_args['source_language']}/{source_domain}/{cfg.train_args['train_filename']}"
            dataset = load_labelled(filepath)
        dataset_embedding = extract_embeddings(cfg, dataset, tokenizer, model)
        embeddings_dict[source_domain] = dataset_embedding
    return embeddings_dict

def create_val_test_embeddings(
    cfg: UFDArguments, tokenizer: UFDTokenizer, model: UFDEmbeddingModel, dataset_type: str
) -> Dict:
    """Helper function to generate validation dataset for supervised and unsupervised training.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file
        tokenizer (UFDTokenizer): UFD tokenizer class instance
        model (UFDEmbeddingModel): UFD embedding model class instance

    Returns:
        Dict: dictionary of dataset embeddings for supervised and unsupervised dataset
    """
    embeddings_dict = {}
    target_languages_list = cfg.train_args["target_languages"]
    target_domains_list = cfg.train_args["target_domains"]
    filename = (
        cfg.train_args["val_filename"]
        if dataset_type == "valid"
        else cfg.eval_args["test_filename"]
    )
    for target_language in target_languages_list:
        embeddings_dict[target_language] = {}
        for target_domain in target_domains_list:
            filepath = (
                f"{cfg.data_folder}/{target_language}/{target_domain}/{filename}"
            )
            dataset = load_labelled(filepath)
            dataset_embedding = extract_embeddings(cfg, dataset, tokenizer, model)
            embeddings_dict[target_language][target_domain] = dataset_embedding
    return embeddings_dict

def create_dataset_embedding(cfg: UFDArguments, dataset_type: str) -> Dict:
    """Main helper wrapper function to generate datasets.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file
        dataset_type (str): type of dataset to generate embeddings

    Raises:
        Exception: raise exception for unknown dataset types

    Returns:
        Dict: return dictionary of dataset embeddings.
    """
    device = torch.device(cfg.device)
    config = UFDEmbeddingConfig.from_pretrained(
        cfg.train_args["embedding_model_name"]
    )
    model = UFDEmbeddingModel.from_pretrained(
        cfg.train_args["embedding_model_name"], config=config
    ).to(device)
    tokenizer = UFDTokenizer.from_pretrained(cfg.train_args["embedding_model_name"])

    if dataset_type == "train":
        dataset_embedding_dict = create_train_embeddings(
            cfg,
            tokenizer,
            model,
        )
    elif dataset_type in ["test", "valid"]:
        dataset_embedding_dict = create_val_test_embeddings(
            cfg, tokenizer, model, dataset_type
        )
    else:
        raise Exception(
            "Invalid value for dataset_type, dataset_type should be train/ valid/ test"
        )

    return dataset_embedding_dict

def create_classifiers(cfg: UFDArguments) -> Dict:
    """Helper function to generate all classifier models, criterion and optimizer.
        One set of classifier, criterion and optimizer required per source domain.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file

    Returns:
        Dict: return dictionary of models, criterion and optimizer group by source domain.
    """
    device = torch.device(cfg.device)
    config = UFDClassifierConfig(
        embed_dim=cfg.train_args["out_dim"],
        num_class=cfg.train_args["num_class"],
        initrange=cfg.train_args["initrange"],
    )
    classifiers = {}
    for domain in cfg.train_args["source_domains"]:
        classifiers[domain] = {}
        classifiers[domain]["model"] = UFDClassifierModel(config).to(device)
        classifiers[domain]["criterion"] = nn.CrossEntropyLoss().to(device)
        classifiers[domain]["optimizer"] = Adam(
            classifiers[domain]["model"].parameters(),
            lr=cfg.train_args["learning_rate"],
        )
    return classifiers

def get_source2target_domain_mapping(source_domains: List[str], target_domains: List[str]) -> Dict[str, List[str]]:
    """Helper function to return cross domains keys for each domains.

    Args:
        source_domains (List[str]): list of all source domains
        target_domains (List[str]): list of all target domains

    Returns:
        Dict[str, List[str]]: return a list of cross domains per domain keys.
    """
    mapping = {}
    for dom in source_domains:
        mapping[dom] = [d for d in target_domains if d != dom]
    return mapping

def generate_train_val_dataset(cfg: UFDArguments) -> Tuple[Dict[str, float]]:
    """Helper function to generate train and validation datasets.
        Load dataset object from cache if available, else call the dataset embeddings creation methods.

    Args:
        cfg (UFDArguments): UFDArguments config load from configuration file

    Returns:
        Tuple(Dict[str, float]): return the generated train and validation dictionaries.
    """
    train_cache_path = str(
        pathlib.Path(cfg.cache_folder).joinpath(
            cfg.train_args["train_cache_filename"]
        )
    )
    if os.path.isfile(train_cache_path):
        with open(train_cache_path, "rb") as handle:
            train_data = pickle.load(handle)
    else:
        train_data = create_dataset_embedding(cfg, dataset_type="train")
        with open(train_cache_path, "wb") as handle:
            pickle.dump(train_data, handle)

    valid_cache_path = str(
        pathlib.Path(cfg.cache_folder).joinpath(
            cfg.train_args["val_cache_filename"]
        )
    )
    if os.path.isfile(valid_cache_path):
        with open(valid_cache_path, "rb") as handle:
            valid_data = pickle.load(handle)
    else:
        valid_data = create_dataset_embedding(cfg, dataset_type="valid")
        with open(valid_cache_path, "wb") as handle:
            pickle.dump(valid_data, handle)

    return train_data, valid_data

def parse_args_and_load_config(config_path: str = "config/ufd_config.json"):
    """Args parser helper method

    Returns:
        UFDArguments: UFDArguments instance with parsed args
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default=config_path)
    args = parser.parse_args()
    with open(args.cfg_path, "r") as cfg_file:
        cfg = json.load(cfg_file)
    ufd_args = UFDArguments(**cfg)
    return ufd_args


# -*- coding:utf8 -*-
"""
This file is part of the repo: https://github.com/tencent-ailab/hifi3dface

If you find the code useful, please cite our paper:

"High-Fidelity 3D Digital Human Head Creation from RGB-D Selfies."
ACM Transactions on Graphics 2021
Code: https://github.com/tencent-ailab/hifi3dface

Copyright (c) [2020-2021] [Tencent AI Lab]

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

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

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

import numpy as np
import cv2
import math
from .PoseTools import PoseTools

class chooseFrame(object):
    def __init__(self):
        pass

    @staticmethod
    def check_eye_close(pts2d, flag=0):
        # in : n * 2
        left_eye_lmk_index_v = [
            58 - 1,
            59 - 1,
            53 - 1,
            57 - 1,
            54 - 1,
            56 - 1,
        ]  # left eye lmk, vertical index
        left_eye_lmk_index_h = [52 - 1, 55 - 1]  # left eye lmk, horizontal index

        right_eye_lmk_index_v = [
            66 - 1,
            67 - 1,
            61 - 1,
            65 - 1,
            62 - 1,
            64 - 1,
        ]  # right eye lmk, vertical index
        right_eye_lmk_index_h = [60 - 1, 63 - 1]  # right eye lmk, horizontal index
        # flag=0->mid pose, flag=1->left/right pose
        if flag == 0:
            thred = 0.2
        else:
            thred = 0.5
        left_eye_distance = 0
        right_eye_distance = 0
        for ind in range(0, len(left_eye_lmk_index_v), 2):
            x1 = pts2d[left_eye_lmk_index_v[ind]][0]
            y1 = pts2d[left_eye_lmk_index_v[ind]][1]
            x2 = pts2d[left_eye_lmk_index_v[ind + 1]][0]
            y2 = pts2d[left_eye_lmk_index_v[ind + 1]][1]

            left_eye_distance = left_eye_distance + np.linalg.norm(
                np.array([x1, y1]) - np.array([x2, y2])
            )

            x3 = pts2d[right_eye_lmk_index_v[ind]][0]
            y3 = pts2d[right_eye_lmk_index_v[ind]][1]
            x4 = pts2d[right_eye_lmk_index_v[ind + 1]][0]
            y4 = pts2d[right_eye_lmk_index_v[ind + 1]][1]

            right_eye_distance = right_eye_distance + np.linalg.norm(
                np.array([x3, y3]) - np.array([x4, y4])
            )
        x5 = pts2d[left_eye_lmk_index_h[0]][0]
        y5 = pts2d[left_eye_lmk_index_h[0]][1]
        x6 = pts2d[left_eye_lmk_index_h[1]][0]
        y6 = pts2d[left_eye_lmk_index_h[1]][1]
        left_eyes_ratio = left_eye_distance / np.linalg.norm(
            np.array([x5, y5]) - np.array([x6, y6])
        )

        x7 = pts2d[right_eye_lmk_index_h[0]][0]
        y7 = pts2d[right_eye_lmk_index_h[0]][1]
        x8 = pts2d[right_eye_lmk_index_h[1]][0]
        y8 = pts2d[right_eye_lmk_index_h[1]][1]
        right_eyes_ratio = right_eye_distance / np.linalg.norm(
            np.array([x7, y7]) - np.array([x8, y8])
        )
        # print("left_eyes_ratio right_eyes_ratio ",left_eyes_ratio,right_eyes_ratio)
        result_status = 1
        if left_eyes_ratio < thred or right_eyes_ratio < thred:
            result_status = result_status - 2
        return result_status

    @staticmethod
    def get_abs_angle_by_orth_pnp(pt2d_all):
        #################   pt2d: (m*2*n),m is the number of pics ,n is the number of landmarks ############

        preset_PT_3d = np.array(
            [
                [
                    -68.1756973070174,
                    -68.1146470361586,
                    -66.0249619011911,
                    -62.5428837026772,
                    -56.7464771621017,
                    -47.0523705994082,
                    -31.2188583709116,
                    -15.3879495971038,
                    -4.08401359484161,
                    16.5929626641805,
                    33.7166876786039,
                    48.8605803831152,
                    56.8829519914283,
                    62.4308036010388,
                    65.3272669641152,
                    67.3733798672488,
                    67.1870150203456,
                    -54.3237949408625,
                    -45.9681810043115,
                    -35.0726032124679,
                    -24.8322654940756,
                    -12.7729299752417,
                    -45.8693519937360,
                    -36.3826769033241,
                    -25.9141249360281,
                    -14.8096785082511,
                    10.7624649776308,
                    20.5206483307513,
                    33.1056897313524,
                    46.6258909064902,
                    54.0926809905806,
                    12.6266583822557,
                    22.7049645972830,
                    34.4879223577708,
                    44.0535131917368,
                    -41.2673480094191,
                    -35.6089130191105,
                    -21.7764329890458,
                    -16.7468454026988,
                    -25.5757409620726,
                    -37.2345564965293,
                    -28.3139941581866,
                    -31.6768869306443,
                    17.6041471849675,
                    22.6210777240253,
                    34.0961319432436,
                    41.9045520975750,
                    36.7447958534758,
                    25.3302610383278,
                    27.9211579237679,
                    31.9237846382224,
                    -0.583686964484780,
                    0.144336413802457,
                    0.269583123001013,
                    0.229647359245878,
                    -10.8552873508045,
                    8.80169774995196,
                    -14.1943604935433,
                    12.3636086468133,
                    -17.6260973432672,
                    16.1859617393272,
                    -11.8561671740247,
                    -6.14033382730938,
                    0.122603013080397,
                    4.51988955136432,
                    11.2784146784043,
                    -22.3665467791695,
                    -5.12694742849893,
                    -0.261962445311256,
                    5.04729862351575,
                    21.9230047599515,
                    2.19959699419832,
                    -14.3228422073861,
                    14.7897345459004,
                    16.5511881096132,
                    -16.7100684305587,
                    -18.5092145428952,
                    19.6847364244945,
                    -10.4805964476244,
                    0.189005281495954,
                    10.7239229159183,
                    -10.1233814192337,
                    0.733705124663359,
                    10.7879717245791,
                    -7.09363186602643,
                    9.83684660832385,
                ],
                [
                    53.3107940578386,
                    33.6623335532468,
                    13.3328235939212,
                    -6.70656176326401,
                    -26.2279897952572,
                    -39.8772599748791,
                    -50.6031722705214,
                    -55.0682077751379,
                    -54.7205192995213,
                    -54.4383527884069,
                    -45.6074643123763,
                    -36.7972070532683,
                    -26.7844743779424,
                    -8.31142380648254,
                    12.9211104101711,
                    32.3868187022165,
                    52.8705611396547,
                    55.7094497224889,
                    65.4640357632679,
                    68.0262305498897,
                    67.0594026704289,
                    65.1248359510029,
                    58.7888308171010,
                    60.8175914739616,
                    60.2800548576361,
                    57.9635621511934,
                    64.1839275608599,
                    66.9301273508612,
                    68.1410948144991,
                    66.2383559145425,
                    55.8636845216315,
                    57.9597060181801,
                    62.0114248050252,
                    61.3778094800454,
                    59.3384362324800,
                    45.4503766622321,
                    49.4789679497268,
                    49.6233682727380,
                    44.7218523846120,
                    42.5597018661359,
                    43.5828799351429,
                    50.9429154037993,
                    41.9671755981695,
                    44.3637555869468,
                    49.5010235719469,
                    50.2143139274464,
                    46.1853031649069,
                    43.5934880814510,
                    43.1457356201204,
                    50.7859006873030,
                    41.9688807138128,
                    52.2133605010753,
                    40.8606102496554,
                    29.8673390697036,
                    18.5682308483948,
                    47.8126617850450,
                    46.9061150639963,
                    23.1271685381814,
                    23.8289873865312,
                    13.2948429656895,
                    12.6521944496779,
                    10.1998621404392,
                    8.07936564185550,
                    7.02048790726341,
                    7.96680610142852,
                    8.85964829325156,
                    -15.4685231851887,
                    -5.83694476072372,
                    -7.18510079110145,
                    -6.04296059535911,
                    -14.6904905468550,
                    -21.9680454218847,
                    -9.29459523117953,
                    -9.09722248628968,
                    -20.0443614430531,
                    -20.5127768888248,
                    -15.2511924645571,
                    -14.5447326599110,
                    -14.2644213165510,
                    -14.4990315246071,
                    -14.0522679543919,
                    -14.6584850106918,
                    -14.8322370025323,
                    -14.3013495881650,
                    -22.1447253850458,
                    -21.1082651587302,
                ],
                [
                    -23.8053118343091,
                    -27.0445038733228,
                    -25.4519279993035,
                    -21.3747758708710,
                    -10.0427175318756,
                    2.26734761619467,
                    18.4816765410735,
                    34.3221442669071,
                    40.5328508971239,
                    34.8712700765862,
                    22.2864591376731,
                    6.21563701787038,
                    -13.1134717329451,
                    -16.4644090354300,
                    -23.5839491111222,
                    -20.1740307731121,
                    -20.9111590403862,
                    14.2565903252940,
                    24.3935737144901,
                    32.3840414727001,
                    37.1291429990255,
                    39.9932199086069,
                    27.2161675645801,
                    33.7749032608401,
                    36.8089431088900,
                    37.8626154364139,
                    40.2974535831101,
                    38.6138476948733,
                    33.2010273723652,
                    22.5020632801053,
                    13.7408293547814,
                    38.2504652330102,
                    37.6996568238641,
                    34.4917209315425,
                    28.6782251624995,
                    25.2856914494286,
                    31.2533888265002,
                    32.6627868245544,
                    30.0533047871161,
                    32.5050230824477,
                    28.5422698190148,
                    33.5168371456454,
                    31.5280683519493,
                    31.0209067620750,
                    33.3543299414210,
                    32.2212854309802,
                    24.3953248646461,
                    28.9265332073268,
                    33.0077738148317,
                    33.6859149016717,
                    31.0602176320082,
                    40.3393280766325,
                    45.6002020387433,
                    55.0202454254639,
                    62.8978452597018,
                    34.8476961173754,
                    36.5202392654199,
                    43.3278891974560,
                    45.9282464018630,
                    43.4287421024055,
                    47.8949929648990,
                    48.6608554911977,
                    49.5354062065318,
                    53.6401343965816,
                    50.5486291847995,
                    47.9518469656950,
                    40.8297273633096,
                    53.0054425832235,
                    54.0679016180019,
                    53.2652104521736,
                    41.4592065997822,
                    51.9548140335428,
                    48.7692056980377,
                    48.8730272580562,
                    44.6779594870922,
                    44.0230121149454,
                    43.2977936185804,
                    42.8604352591731,
                    48.2764306655049,
                    51.5130563672583,
                    48.6159735298573,
                    47.5215996627314,
                    50.4681746449472,
                    47.7544272187955,
                    50.4640358517105,
                    49.7103891466216,
                ],
            ]
        )

        preset_PT_3d = preset_PT_3d.transpose()  # n * 3

        length = len(pt2d_all)
        phi_all = []
        gamma_all = []
        for i in range(length):

            one_pt2d = pt2d_all[i].transpose()  # n * 2
            pt2d = one_pt2d - np.tile(np.mean(one_pt2d, axis=0), (one_pt2d.shape[0], 1))
            phi, gamma, theta, t3d, f = PoseTools.pnp_orth(pt2d, preset_PT_3d)
            phi1 = phi * 180 / 3.14

            if phi1 > 90:
                phi1 = phi1 - 180
            if phi1 < -90:
                phi1 = phi1 + 180
            phi_all.append(-phi1)
            gamma_all.append(-gamma * 180 / 3.14)

        return np.array(phi_all), np.array(gamma_all)

    @staticmethod
    def ransac_rigid_trans(source, target, n_pairs):
        # in : n * 3
        allowed_error = 6
        nBest = 0
        nPoints = source.shape[0]
        inliers = []

        def computeInliers(trans, source, target, allowedError):
            allowed_error = 6
            nPoints = source.shape[0]

            idx = []
            for i in range(nPoints):
                p1 = source[i, :]
                p2 = target[i, :]
                p1 = np.hstack((p1, 1))
                p1Trans = np.dot(trans, p1.transpose())
                distance = np.linalg.norm(p1Trans - p2)
                if distance < allowed_error:
                    idx.append(i)
            return idx

        for i in range(50):
            samples = np.random.randint(0, nPoints, size=10)

            currentTransformation = PoseTools.fit_icp_RT_no_scale(
                source[samples, :], target[samples, :]
            )
            currentInliers = computeInliers(
                currentTransformation, source, target, allowed_error
            )

            if not currentInliers:
                continue
            n_current_inliers = len(currentInliers)
            # print(n_current_inliers)
            if n_current_inliers > nBest:
                nBest = n_current_inliers
                inliers = currentInliers
        if nBest <= n_pairs:
            trans = []
        trans = PoseTools.fit_icp_RT_no_scale(source[inliers, :], target[inliers, :])
        return trans, inliers, nBest

    @staticmethod
    def call_one_pose_by_DLT(useful_points, in_ref3d, in_pt3d):
        # in : n * 3
        target = in_ref3d[useful_points, :]
        source = in_pt3d[useful_points, :]

        valid_source = []
        valid_target = []
        for i in range(len(target)):
            if target[i][2] > 100 and target[i][2] < 1000:
                valid_target.append(i)
            if source[i][2] > 100 and source[i][2] < 1000:
                valid_source.append(i)
        valid_target = set(valid_target)
        valid_source = set(valid_source)
        valid_ind = list(valid_target & valid_source)

        if len(valid_ind) < 15:
            return -1, None
        source = source[valid_ind, :]
        target = target[valid_ind, :]

        trans, inliers, nInliners = chooseFrame.ransac_rigid_trans(source, target, 10)
        # print(trans, len(inliers), nInliners)
        if nInliners <= 10:
            return -1, None
        return 1, trans

    @staticmethod
    def find_close_angle_index(points, angle):
        index = []
        state = points >= angle
        for i in range(len(state) - 1):
            if state[i] != state[i + 1]:
                if abs(points[i] - angle) < abs(points[i + 1] - angle):
                    index.append(i)
                else:
                    index.append(i + 1)
        return index

    @staticmethod
    def rotationMatrixToEulerAngles(R):
        r1 = math.atan2(R[1][0], R[0][0])
        r2 = math.asin(-R[2][0])
        r3 = math.atan2(R[2][1], R[2][2])
        r = [r1, r2, r3]
        return r

    @staticmethod
    def call_motion_blur(img):
        img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        mask = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
        num_depth = cv2.filter2D(img_gray, -1, mask)
        # num_depth = scipy.ndimage.correlate(img_gray, mask,mode='nearest')
        (mean, stddv) = cv2.meanStdDev(num_depth)
        score = stddv[0][0]
        return score

    @staticmethod
    def find_chin(pt3d_all, in_up_down_all, left_right_all, ref_id, img_all):
        id_chin = []
        part_chin = []
        flag_ok = 0
        ref_angle = in_up_down_all[ref_id]
        up_down_all = in_up_down_all
        up_down_all[np.where(abs(left_right_all) > 20)] = ref_angle
        min_chin = np.min(np.where(up_down_all == np.min(up_down_all)))
        thres = np.max((7, (ref_angle - up_down_all[min_chin]) / 3 + ref_angle))
        preset_max_chin_angle = ref_angle - thres

        if up_down_all[min_chin] > preset_max_chin_angle:
            flag_ok = -1
            trans_chin = []

        chin_cadidate = chooseFrame.find_close_angle_index(
            up_down_all, preset_max_chin_angle
        )
        if len(chin_cadidate) == 1:
            idx_begin = chin_cadidate[0]
            idx_end = len(up_down_all)
        else:
            idx_begin = chin_cadidate[-2]
            idx_end = chin_cadidate[-1]
        useful_points = [
            3,
            4,
            5,
            6,
            7,
            8,
            9,
            10,
            11,
            12,
            13,
            17,
            22,
            23,
            24,
            25,
            30,
            31,
            32,
            33,
            34,
            66,
            70,
            71,
            74,
            75,
            84,
            85,
        ]
        ref3d = pt3d_all[ref_id][:, useful_points]
        valid_target = np.where(
            (ref3d[2, :] > 100) & (ref3d[2, :] < 1000)
        )  # Depth prior of face selfie scene
        blur_score_all = []
        angles_all = []
        idx_all = []
        trans_all = []
        count_ok_frame = 0
        for i in range(idx_begin, idx_end):
            pt3d = pt3d_all[i][:, useful_points]
            source = pt3d
            target = ref3d
            valid_source = np.where(
                (source[2, :] > 100) & (source[2, :] < 1000)
            )  # Depth prior of face selfie scene

            valid_source1 = []
            for j in range(len(valid_source[0])):
                valid_source1.append(valid_source[0][j])

            valid_target1 = []
            for j in range(len(valid_target[0])):
                valid_target1.append(valid_target[0][j])
            valid_ind = list(set(valid_source1).intersection(set(valid_target1)))
            if len(valid_ind) < 15:
                print("not enough inliner in pose calulation", i)
            source = source[:, valid_ind]
            target = target[:, valid_ind]
            trans, inliers, nInliners = chooseFrame.ransac_rigid_trans(
                source.transpose(), target.transpose(), 10
            )
            if nInliners <= 10:
                print(
                    "not enough inliner in pose calulation  ",
                    i,
                    " nInliners: ",
                    nInliners,
                )
            angle = chooseFrame.rotationMatrixToEulerAngles(trans[:, 0:3])
            angle1 = ref_angle + -1 * abs(angle[2] * 180 / 3.14)
            if abs(angle1) < 30:
                b = chooseFrame.call_motion_blur(img_all[i])
                blur_score_all.append(b)
                angles_all.append(angle1)
                idx_all.append(i)
                trans_all.append(trans)
                count_ok_frame = count_ok_frame + 1

        if count_ok_frame == 0:
            print("no chin   !!!!")
            flag_ok = -1
            return [], [], [], flag_ok

        threshold = -20
        for i in range(len(angles_all)):
            angles_all[i] = angles_all[i] - threshold
        b = np.argsort((abs(np.array(angles_all))))
        a = np.sort(abs(np.array(angles_all)))
        len1 = int(np.ceil(len(b) / 3))
        c = angles_all.index(max(angles_all[0 : len1 + 1]))
        id1 = b[c]
        id_chin = idx_all[id1]
        trans_chin = trans_all[id1]
        part_chin = chin_cadidate
        flag_ok = 1
        return id_chin, trans_chin, part_chin, flag_ok

    @staticmethod
    def find_two_part_candidate_by_sort(left_right_all, ref_id, part_chin):
        maxx = -1000000
        minn = 1000000
        leftmost_ind = -1
        rightmost_ind = -1
        for i in range(len(left_right_all)):
            if maxx < left_right_all[i]:
                maxx = left_right_all[i]
                leftmost_ind = i
            if minn >= left_right_all[i]:
                minn = left_right_all[i]
                rightmost_ind = i
        Left_min_ANGEL = 10
        Right_min_ANGEL = -10
        Left_MAX_ANGEL = 50
        Right_MAX_ANGEL = -50
        if (left_right_all[ref_id]) > (Left_min_ANGEL) - 5:
            Left_min_ANGEL = left_right_all[ref_id] + 5
        if left_right_all[ref_id] < Right_min_ANGEL + 5:
            Right_min_ANGEL = left_right_all[ref_id] - 5
        if left_right_all[leftmost_ind] < Left_MAX_ANGEL:
            Left_MAX_ANGEL = left_right_all[leftmost_ind] - 1
        if left_right_all[rightmost_ind] > Right_MAX_ANGEL:
            Right_MAX_ANGEL = left_right_all[rightmost_ind] + 1
        print(
            "   min left is : ", Left_min_ANGEL, "    max left  is : ", Left_MAX_ANGEL
        )
        print(
            "   min right is : ",
            Right_min_ANGEL,
            "    max right is : ",
            Right_MAX_ANGEL,
        )

        left_with_chin = np.where(
            (left_right_all > Left_min_ANGEL) & (left_right_all < Left_MAX_ANGEL)
        )
        right_with_chin = np.where(
            (left_right_all < Right_min_ANGEL) & (left_right_all > Right_MAX_ANGEL)
        )
        left1 = np.sort(
            list(set(list(left_with_chin[0])).difference(set(list(part_chin))))
        )
        right1 = np.sort(
            list(set(list(right_with_chin[0])).difference(set(list(part_chin))))
        )

        part_candidate_left = left1
        if len(part_candidate_left) < 1:
            part_candidate_left = left_with_chin
            print(" warning！！！ left face is divided into the chin section")
        part_candidate_right = right1
        if len(part_candidate_right) < 1:
            part_candidate_right = right_with_chin
            print(" warning！！！ right face is divided into the chin section")
        return part_candidate_left, part_candidate_right

    @staticmethod
    def find_best_imgs_in_part(
        ref_lr_angle,
        part_candidate,
        img_all,
        pt3d_all,
        pt2d_all,
        ref_id,
        part_lmk_index,
    ):
        """
        Between the two segments, limit the minimum interval angle to sample, at the same time delete the points with obviously poor quality,
        and recalculate the exact pose relative to the first frame
        Here are the calculations of the rotation and translation relative to the first frame, then use 3d points as icp, plus ransac to filter
        The previous pnp method cannot be used here, because the average 3d point of fwh does not match this person and cannot be used to filter the outliner

        The idea here is almost the same as choosing the chin, the difference is that only one-sided inliner is used when calculating the pose,
        and the logic of selecting the last few pictures
        """
        useful_points = part_lmk_index
        ref3d = pt3d_all[ref_id][:, useful_points]
        valid_target = np.where(
            (ref3d[2, :] > 100) & (ref3d[2, :] < 1000)
        )  # Depth prior of face selfie scene
        angles_all = []
        blur_score_all = []
        good_part_candidate = []
        close_eye_all = []
        trans_all = []
        count_ok_frame = 0
        for ite in range(len(part_candidate)):
            i = part_candidate[ite]
            pt3d = pt3d_all[i][:, useful_points]
            source = pt3d
            target = ref3d
            valid_source = np.where((source[2, :] > 100) & (source[2, :] < 1000))
            valid_ind = list(
                set(list(valid_source[0])).intersection(set(list(valid_target[0])))
            )
            if len(valid_ind) < 15:
                print("not enough inliner in pose calulation", i)
            source = source[:, valid_ind]
            target = target[:, valid_ind]
            trans, inliers, nInliners = chooseFrame.ransac_rigid_trans(
                source.transpose(), target.transpose(), 10
            )
            if nInliners <= 10:
                print(
                    "not enough inliner in pose calulation  ",
                    i,
                    " nInliners: ",
                    nInliners,
                )
            angle = chooseFrame.rotationMatrixToEulerAngles(trans[:, 0:3])
            angle = ref_lr_angle + angle[1] * 180 / 3.14
            if abs(angle) < 40 and abs(angle) > 10:
                angles_all.append(angle)
                good_part_candidate.append(i)
                close_eye_all.append(
                    chooseFrame.check_eye_close(pt2d_all[i].transpose(), 1)
                )
                count_ok_frame = count_ok_frame + 1
                trans_all.append(trans)
                score = chooseFrame.call_motion_blur(img_all[i])
                blur_score_all.append(score)
        if len(angles_all) == 0:
            out_idx = 0
            out_trans_all = []
            return
        angles_all = np.array(angles_all)
        angles_all = abs(angles_all)
        out_idx = []
        num_class = 7
        range_angle = 1 + max(abs(angles_all)) - min(abs(angles_all))
        per_angle = range_angle / (num_class + 1)
        class1 = np.floor((abs(angles_all) - min(abs(angles_all))) / per_angle)
        index_in_all = []
        final_motion_blur_score = []
        for classid in range(int(max(class1)), 0, -1):
            caditate = np.where(class1 == classid)
            ok_cadidate = np.array(caditate[0])[
                np.where(np.array(close_eye_all)[caditate] > 0)
            ]
            if len(ok_cadidate) > 0:
                indx = np.argmax(np.array(blur_score_all)[ok_cadidate])
                final_motion_blur_score.append(blur_score_all[int(ok_cadidate[indx])])
                out_idx.append(good_part_candidate[ok_cadidate[indx]])
                index_in_all.append(ok_cadidate[indx])
        if len(out_idx) == 0:
            out_idx = np.argmax(angles_all)
            out_trans_all = trans_all[out_idx]
            out_score = 1
        else:
            out_trans_all = np.zeros((len(out_idx), 3, 4))
            for i in range(len(out_idx)):
                out_trans_all[i] = trans_all[index_in_all[i]]
            out_score = final_motion_blur_score
        return out_idx, out_trans_all, out_score

    @staticmethod
    def find_left_right_part(
        left_right_all, img_all, pt3d_all, pt2d_all, ref_id, part_chin
    ):
        flag_ok = 1
        (
            part_candidate_left,
            part_candidate_right,
        ) = chooseFrame.find_two_part_candidate_by_sort(
            left_right_all, ref_id, part_chin
        )
        if len(part_candidate_left) == 0:
            print(" error！！！ Missing left face data")
            flag_ok = -1
        if len(part_candidate_right) == 0:
            print("  error！！！ Missing right face data")
            flag_ok = -1
        left_lmk_index = [
            18,
            19,
            20,
            21,
            22,
            23,
            24,
            25,
            26,
            36,
            37,
            38,
            39,
            40,
            41,
            42,
            43,
            56,
            58,
            60,
            62,
            63,
            67,
            68,
            73,
            76,
            77,
            79,
            82,
            85,
        ]
        right_lmk_index = [
            27,
            28,
            29,
            30,
            31,
            32,
            33,
            34,
            35,
            44,
            45,
            46,
            47,
            48,
            49,
            50,
            51,
            57,
            59,
            61,
            65,
            66,
            70,
            71,
            74,
            75,
            78,
            81,
            84,
            86,
        ]
        for i in range(len(left_lmk_index)):
            left_lmk_index[i] = left_lmk_index[i] - 1
            right_lmk_index[i] = right_lmk_index[i] - 1
        idx_left, trans_left, score_left = chooseFrame.find_best_imgs_in_part(
            left_right_all[ref_id],
            part_candidate_left,
            img_all,
            pt3d_all,
            pt2d_all,
            ref_id,
            left_lmk_index,
        )
        idx_right, trans_right, score_right = chooseFrame.find_best_imgs_in_part(
            left_right_all[ref_id],
            part_candidate_right,
            img_all,
            pt3d_all,
            pt2d_all,
            ref_id,
            right_lmk_index,
        )
        if len(trans_left) == 0:
            print(" error!!! some wrong in left pose calculation")
            flag_ok = -1
        if len(trans_right) == 0:
            print("error!!! some wrong in right pose calculation")
            flag_ok = -1
        return (
            idx_left,
            trans_left,
            idx_right,
            trans_right,
            flag_ok,
            score_left,
            score_right,
        )


import numpy as np
import os

# import time
import copy
import wave
import wget

# from omegaconf import OmegaConf

import matplotlib.pyplot as plt
import librosa.display

import IPython.display as ipd

# import pyaudio as pa
import librosa

# import nemo
import nemo.collections.asr as nemo_asr
from nemo.core.classes import IterableDataset
from nemo.core.neural_types import NeuralType, AudioSignal, LengthsType
import torch
from torch.utils.data import DataLoader

# sample rate, Hz
SAMPLE_RATE = 16000
# import pdb; pdb.set_trace()

# vad_model = nemo_asr.models.EncDecClassificationModel.from_pretrained(
#     "vad_marblenet"
# )
# vad_model = nemo_asr.models.EncDecClassificationModel.from_pretrained(
#     model_name="MarbleNet-3x2x64-Telephony"
# )
vad_model = nemo_asr.models.EncDecClassificationModel.restore_from(
    "/home/malar/work/test/vad_telephony_marblenet.nemo"
)
# vad_model = nemo_asr.models.EncDecClassificationModel.from_pretrained(
#     model_name="vad_telephony_marblenet"
# )
# Preserve a copy of the full config
cfg = copy.deepcopy(vad_model._cfg)
# print(OmegaConf.to_yaml(cfg))

vad_model.preprocessor = vad_model.from_config_dict(cfg.preprocessor)

# Set model to inference mode
vad_model.eval()
vad_model = vad_model.to(vad_model.device)
# import pdb; pdb.set_trace()

# simple data layer to pass audio signal
class AudioDataLayer(IterableDataset):
    @property
    def output_types(self):
        return {
            "audio_signal": NeuralType(
                ("B", "T"), AudioSignal(freq=self._sample_rate)
            ),
            "a_sig_length": NeuralType(tuple("B"), LengthsType()),
        }

    def __init__(self, sample_rate):
        super().__init__()
        self._sample_rate = sample_rate
        self.output = True

    def __iter__(self):
        return self

    def __next__(self):
        if not self.output:
            raise StopIteration
        self.output = False
        return (
            torch.as_tensor(self.signal, dtype=torch.float32),
            torch.as_tensor(self.signal_shape, dtype=torch.int64),
        )

    def set_signal(self, signal):
        self.signal = signal.astype(np.float32) / 32768.0
        self.signal_shape = self.signal.size
        self.output = True

    def __len__(self):
        return 1

data_layer = AudioDataLayer(sample_rate=cfg.train_ds.sample_rate)
data_loader = DataLoader(
    data_layer, batch_size=1, collate_fn=data_layer.collate_fn
)

# inference method for audio signal (single instance)
def infer_signal(model, signal):
    data_layer.set_signal(signal)
    batch = next(iter(data_loader))
    audio_signal, audio_signal_len = batch
    audio_signal, audio_signal_len = (
        audio_signal.to(vad_model.device),
        audio_signal_len.to(vad_model.device),
    )
    logits = model.forward(
        input_signal=audio_signal, input_signal_length=audio_signal_len
    )
    return logits

# class for streaming frame-based VAD
# 1) use reset() method to reset FrameVAD's state
# 2) call transcribe(frame) to do VAD on
#    contiguous signal's frames
class FrameVAD:
    def __init__(
        self,
        model_definition,
        threshold=0.5,
        frame_len=2,
        frame_overlap=2.5,
        offset=10,
    ):
        """
        Args:
          threshold: If prob of speech is larger than threshold, classify the segment to be speech.
          frame_len: frame's duration, seconds
          frame_overlap: duration of overlaps before and after current frame, seconds
          offset: number of symbols to drop for smooth streaming
        """
        self.vocab = list(model_definition["labels"])
        self.vocab.append("_")

        self.sr = model_definition["sample_rate"]
        self.threshold = threshold
        self.frame_len = frame_len
        self.n_frame_len = int(frame_len * self.sr)
        self.frame_overlap = frame_overlap
        self.n_frame_overlap = int(frame_overlap * self.sr)
        timestep_duration = model_definition["AudioToMFCCPreprocessor"][
            "window_stride"
        ]
        for block in model_definition["JasperEncoder"]["jasper"]:
            timestep_duration *= block["stride"][0] ** block["repeat"]
        self.buffer = np.zeros(
            shape=2 * self.n_frame_overlap + self.n_frame_len, dtype=np.float32
        )
        self.offset = offset
        self.reset()

    def _decode(self, frame, offset=0):
        assert len(frame) == self.n_frame_len
        self.buffer[: -self.n_frame_len] = self.buffer[self.n_frame_len :]
        self.buffer[-self.n_frame_len :] = frame
        logits = infer_signal(vad_model, self.buffer).cpu().numpy()[0]
        decoded = self._greedy_decoder(self.threshold, logits, self.vocab)
        return decoded

    @torch.no_grad()
    def transcribe(self, frame=None):
        if frame is None:
            frame = np.zeros(shape=self.n_frame_len, dtype=np.float32)
        if len(frame) < self.n_frame_len:
            frame = np.pad(
                frame, [0, self.n_frame_len - len(frame)], "constant"
            )
        unmerged = self._decode(frame, self.offset)
        return unmerged

    def reset(self):
        """
        Reset frame_history and decoder's state
        """
        self.buffer = np.zeros(shape=self.buffer.shape, dtype=np.float32)
        self.prev_char = ""

    @staticmethod
    def _greedy_decoder(threshold, logits, vocab):
        s = []
        if logits.shape[0]:
            probs = torch.softmax(torch.as_tensor(logits), dim=-1)
            probas, _ = torch.max(probs, dim=-1)
            probas_s = probs[1].item()
            preds = 1 if probas_s >= threshold else 0
            s = [
                preds,
                str(vocab[preds]),
                probs[0].item(),
                probs[1].item(),
                str(logits),
            ]
        return s

# WINDOW_SIZE_RANGE = [0.10, 0.15, 0.20, 0.25, 0.30, 0.5, 0.8]
# # STEP_RANGE = [0.01, 0.02, 0.03]
# # WINDOW_SIZE_RANGE = [0.15, 0.20]
# STEP_RANGE = [0.01, 0.02, 0.03]
WINDOW_SIZE_RANGE = [0.15, 0.20, 0.25]
# STEP_RANGE = [0.01, 0.02, 0.03]
# WINDOW_SIZE_RANGE = [0.15, 0.20]
STEP_RANGE = [0.03, 0.05, 0.07, 0.1]
STEP_LIST = [r for t in STEP_RANGE for r in [t]*len(WINDOW_SIZE_RANGE)]
# STEP_LIST
# STEP_LIST = (
#     [0.01] * len(WINDOW_SIZE_RANGE)
#     + [0.02] * len(WINDOW_SIZE_RANGE)
#     + [0.03] * len(WINDOW_SIZE_RANGE)
# )
WINDOW_SIZE_LIST = WINDOW_SIZE_RANGE * len(STEP_RANGE)

def offline_inference(wave_file, STEP=0.025, WINDOW_SIZE=0.5, threshold=0.5):

    FRAME_LEN = STEP  # infer every STEP seconds
    CHANNELS = 1  # number of audio channels (expect mono signal)
    RATE = 16000  # sample rate, Hz

    CHUNK_SIZE = int(FRAME_LEN * RATE)

    vad = FrameVAD(
        model_definition={
            "sample_rate": SAMPLE_RATE,
            "AudioToMFCCPreprocessor": cfg.preprocessor,
            "JasperEncoder": cfg.encoder,
            "labels": cfg.labels,
        },
        threshold=threshold,
        frame_len=FRAME_LEN,
        frame_overlap=(WINDOW_SIZE - FRAME_LEN) / 2,
        offset=0,
    )

    wf = wave.open(wave_file, "rb")
    # p = pa.PyAudio()

    empty_counter = 0

    preds = []
    proba_b = []
    proba_s = []

    # stream = p.open(
    #     format=p.get_format_from_width(wf.getsampwidth()),
    #     channels=CHANNELS,
    #     rate=RATE,
    #     output=True,
    # )

    data = wf.readframes(CHUNK_SIZE)

    while len(data) > 0:

        data = wf.readframes(CHUNK_SIZE)
        signal = np.frombuffer(data, dtype=np.int16)
        result = vad.transcribe(signal)

        preds.append(result[0])
        proba_b.append(result[2])
        proba_s.append(result[3])

        if len(result):
            # print(result, end="\n")
            empty_counter = 3
        elif empty_counter > 0:
            empty_counter -= 1
            # if empty_counter == 0:
            #     print(" ", end="")

    # p.terminate()
    vad.reset()

    return preds, proba_b, proba_s

# demo_wave = "VAD_demo.wav"
# if not os.path.exists(demo_wave):
#     wget.download(
#         "https://dldata-public.s3.us-east-2.amazonaws.com/VAD_demo.wav",
#         demo_wave,
#     )

demo_wave = "WAL-1201-cust.wav"

wave_file = demo_wave

CHANNELS = 1
RATE = 16000
audio, sample_rate = librosa.load(wave_file, sr=RATE)
dur = librosa.get_duration(audio)
print(dur)

threshold = 0.5

results = []
for STEP, WINDOW_SIZE in zip(
    STEP_LIST,
    WINDOW_SIZE_LIST,
):
    print(f"====== STEP is {STEP}s, WINDOW_SIZE is {WINDOW_SIZE}s ====== ")
    preds, proba_b, proba_s = offline_inference(
        wave_file, STEP, WINDOW_SIZE, threshold
    )
    results.append([STEP, WINDOW_SIZE, preds, proba_b, proba_s])

plt.figure(figsize=[20, 3*len(STEP_LIST)])

num = len(results)
for i in range(num):
    len_pred = len(results[i][2])
    FRAME_LEN = results[i][0]
    ax1 = plt.subplot(num + 1, 1, i + 1)

    ax1.plot(np.arange(audio.size) / sample_rate, audio, "b")
    ax1.set_xlim([-0.01, int(dur) + 1])
    ax1.tick_params(axis="y", labelcolor="b")
    ax1.set_ylabel("Signal")
    ax1.set_ylim([-1, 1])

    proba_s = results[i][4]
    pred = [1 if p > threshold else 0 for p in proba_s]
    ax2 = ax1.twinx()
    ax2.plot(
        np.arange(len_pred) / (1 / results[i][0]),
        np.array(pred),
        "r",
        label="pred",
    )
    ax2.plot(
        np.arange(len_pred) / (1 / results[i][0]),
        np.array(proba_s),
        "g--",
        label="speech prob",
    )
    ax2.tick_params(axis="y", labelcolor="r")
    legend = ax2.legend(loc="lower right", shadow=True)
    ax1.set_ylabel("prediction")

    ax2.set_title(f"step {results[i][0]}s, buffer size {results[i][1]}s")
    ax2.set_ylabel("Preds and Probas")

ax = plt.subplot(num + 1, 1, i + 2)
S = librosa.feature.melspectrogram(
    y=audio, sr=sample_rate, n_mels=64, fmax=8000
)
S_dB = librosa.power_to_db(S, ref=np.max)
librosa.display.specshow(
    S_dB, x_axis="time", y_axis="mel", sr=sample_rate, fmax=8000
)
ax.set_title("Mel-frequency spectrogram")
ax.grid()
plt.show()
ipd.Audio(data=audio, rate=sample_rate)


#!/usr/bin/python
# MIT License
import sys
import os
import numpy as np
import subprocess
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from collections import deque

figDec,decAx = plt.subplots()

figDec.set_facecolor('#F2F1F0')
figDec.canvas.set_window_title('NVDEC Activity Monitor')

# For the comparison
decLine, = decAx.plot([],[])

# The line points in x,y list form
cpuy_list = deque([0]*240)
cpux_list = deque(np.linspace(60,0,num=240))

fill_lines_cpu=0

def initDecGraph():
    global decAx
    global decLine
    global fill_lines_cpu

    decAx.set_xlim(60, 0)
    decAx.set_ylim(-5, 105)
    decAx.set_title('NVDEC History')
    decAx.set_ylabel('NVDEC Usage (%)')
    decAx.set_xlabel('Seconds');
    decAx.grid(color='gray', linestyle='dotted', linewidth=1)

    decLine.set_data([],[])
    fill_lines_cpu=decAx.fill_between(decLine.get_xdata(),50,0)

    return [decLine] + [fill_lines_cpu]

def updateDecGraph(frame):
    global fill_lines_cpu
    global cpuy_list
    global cpux_list
    global decLine
    global decAx
    global cpu_user_prev
    global cpu_idle_prev

    dump=subprocess.check_output(["adb shell cat /sys/bus/platform/drivers/nvdec/54480000.nvdec/load"],shell=True).decode("utf-8")

    if dump != "":
        actmon_avg=int(dump.strip())

        cpuy_list.popleft()

        cpuy_list.append(actmon_avg/10)

        decLine.set_data(cpux_list,cpuy_list)

        fill_lines_cpu.remove()
        fill_lines_cpu=decAx.fill_between(cpux_list,0,cpuy_list, facecolor='red', alpha=0.50)

    return [decLine] + [fill_lines_cpu]

decAnimation = FuncAnimation(figDec, updateDecGraph, frames=200,
                    init_func=initDecGraph,  interval=50, blit=True)

plt.show()


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths

import os
import cv2
import numpy as np

import sys
sys.path.append("..")

from opts_pose import opts
from detectors.detector_factory import detector_factory

class Clothes_detector():
    def __init__(self,weights):
        opt = opts().init()
        opt.load_model = weights
        os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
        Detector = detector_factory[opt.task]
        self.detector = Detector(opt)
        self.opt = opt

    def detect(self,img):
        results = self.detector.run(img)['results']

        det=[]

        for bbox in results[1]:
            if bbox[4] > self.opt.vis_thresh:

                points=[]
                for i in range(0,len(bbox)):
                    points.append(int(bbox[i]))

                det.append(points)
        return det

    def get_points(self,img):

        det = self.detect(img)

        sk_out = []

        if (len(det)>0):
            for i in range(0,len(det)):

                bbox = det[i]

                #####################    HEAD    #############################
                head = []
                for j in range(5,15):
                    head.append(bbox[j])

                ###################    SHIRT   #############################
                shirt = []
                for j in range(15,31):
                    shirt.append(bbox[j])

                ###################    PANTS   #############################
                pants = []
                for j in range(27,39):
                    pants.append(bbox[j])

                sk_out.append([head,shirt,pants])

        return sk_out

    def get_bbox(self,img):

        det = self.detect(img)

        bbox_out = []

        if (len(det)>0):
            for i in range(0,len(det)):

                bbox = det[i]
                #####################    HEAD    #############################
                diff = int(np.abs(bbox[11]-bbox[13])*0.6)+ int(np.abs(bbox[12]-bbox[14])*0.3)
                head =  [bbox[11],min(bbox[12],bbox[14])-diff,bbox[13],min(bbox[12],bbox[14])+diff]

                ###################    SHIRT   #############################
                xs,ys=[],[]
                for j in range(15,31):
                    if j%2==0:
                        ys.append(bbox[j])
                    else:
                        xs.append(bbox[j])

                reg=int(np.abs(bbox[15]-bbox[17])*0.2)
                shirt = [np.min(xs)-reg,np.min(ys)-reg,np.max(xs)+reg,np.max(ys)]

                ###################    PANTS   #############################
                xp,yp=[],[]
                for j in range(27,39):
                    if j%2==0:
                        yp.append(bbox[j])
                    else:
                        xp.append(bbox[j])

                pants = [np.min(xp)-reg*2,np.min(yp)-reg,np.max(xp)+reg*2,np.max(yp)]

                bbox_out.append([head,shirt,pants])

        return bbox_out

def main():
    detector = Clothes_detector('../models/multi_pose_dla_3x.pth')

    cam = cv2.VideoCapture(0)
    while True:
        _, img = cam.read()
        cv2.imshow('input', img)

        # det = detector.detect(img)
        # if (len(det)>0):
        #     clothes_img = print_clothes(img,det[0])
        #     cv2.imshow('Clothes', clothes_img)

        bboxes = detector.get_bbox(img)

        clothes_img = img.copy()
        if ( len(bboxes)>0):
            for box in bboxes:
                head, shirt, pants = box[0], box[1], box[2]

                cv2.rectangle(clothes_img,(head[0],head[1]),(head[2],head[3]),(0,0,255),3)
                cv2.rectangle(clothes_img,(shirt[0],shirt[1]),(shirt[2],shirt[3]),(0,255,0),3)
                cv2.rectangle(clothes_img,(pants[0],pants[1]),(pants[2],pants[3]),(255,0,0),3)

        cv2.imshow('Clothes', clothes_img)

        skts = detector.get_points(img)

        if ( len(skts)>0):
            for skt in skts:
                head, shirt, pants = skt[0], skt[1], skt[2]
                #print("sizes: [{},{},{}]".format(len(head),len(shirt),len(pants)))

                for i in range(0,len(head)//2):
                    cv2.circle(img,(head[2*i],head[2*i+1]), 7, (0,0,255), -1)

                for i in range(0,len(shirt)//2):
                    cv2.circle(img,(shirt[2*i],shirt[2*i+1]), 7, (0,255,0), -1)

                for i in range(0,len(pants)//2):
                    cv2.circle(img,(pants[2*i],pants[2*i+1]), 7, (255,0,0), -1)

        cv2.imshow('Poses', img)

        if cv2.waitKey(1) == 27:
            return  # esc to quit

if __name__ == '__main__':
    main()


import cv2
import numpy as np

# Load Yolo
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
classes = []
with open("coco.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# Colours
colors = []
colors.append((255, 0, 0))
colors.append((0, 0, 255))
colors.append((0, 255, 0))
colors.append((255, 255, 0))
colors.append((255, 0, 255))
colors.append((0, 255, 255))
colors.append((102, 0, 0))
colors.append((105, 105, 105))
colors.append((106, 90, 205))
colors.append((75, 0, 130))
colors.append((0, 0, 128))
colors.append((128, 0, 0))
colors.append((255, 215, 0))
colors.append((255, 69, 0))

# Loading image
img = cv2.imread("test1.jpg")

# if image is test 2
# img = cv2.resize(img, None, fx=0.3, fy=0.3)

height, width, channels = img.shape

# Detecting objects
blob = cv2.dnn.blobFromImage(img, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
net.setInput(blob)
outs = net.forward(output_layers)

# Showing information on the screen
class_ids = []
confidences = []
boxes = []
for out in outs:
    for detection in out:
        scores = detection[5:]
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        if confidence > 0.5:
            # Object detected
            center_x = int(detection[0] * width)
            center_y = int(detection[1] * height)
            w = int(detection[2] * width)
            h = int(detection[3] * height)

            # Rectangle coordinates
            x = int(center_x - w / 2)
            y = int(center_y - h / 2)

            boxes.append([x, y, w, h])
            confidences.append(float(confidence))
            class_ids.append(class_id)

indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
# print(indexes)
font = cv2.FONT_HERSHEY_PLAIN
for i in range(len(boxes)):
    if i in indexes:
        x, y, w, h = boxes[i]
        label = str(classes[class_ids[i]])
        cv2.rectangle(img, (x, y), (x + w, y + h), colors[np.random.randint(0, len(colors))], 2)
        cv2.putText(img, label, (x, y + 30), font, 2, colors[np.random.randint(0, len(colors))], 2)

cv2.imshow("Image", img)
cv2.waitKey(0)
cv2.destroyAllWindows()


from gensim.models.coherencemodel import CoherenceModel
from visualization import save_topic_visualization, save_topic_wordclouds
from sklearn.metrics import silhouette_score
from shared.utils import load_from_pickle
from shared.utils import dump_to_pickle
from shared.utils import dump_to_json
from shared.utils import dump_to_txt
from shared.utils import make_dirs
from collections import Counter
from tqdm import tqdm
import pandas as pd
import numpy as np
from tm import TM
import logging
import time
import glob
import os

pd.set_option('display.max_columns', None)
pd.set_option("max_rows", None)

class Evaluation(object):
    """ Class for generating topic model and evaluation files.
    """
    def __init__(self, lang_code, method="BERT", version="1.1", norm_type="LEMMA", pre_trained_name="distilbert-base-nli-mean-tokens",
                 num_words=10, n_neighbors=15, n_components=5, min_dist=0.0, umap_metric='cosine', min_cluster_size=30, cluster_metric='euclidean',
                 cluster_selection_method='eom', prediction_data=True, random_state=42, num_wordcloud_words=150):

        self.lang_code = lang_code
        self.method = method
        self.version = version
        self.norm_type = norm_type
        self.num_words = num_words
        self.pre_trained_name = pre_trained_name
        self.n_neighbors = n_neighbors
        self.n_components = n_components
        self.min_dist = min_dist
        self.umap_metric = umap_metric
        self.min_cluster_size = min_cluster_size
        self.cluster_metric = cluster_metric
        self.cluster_selection_method = cluster_selection_method
        self.prediction_data = prediction_data
        self.random_state = random_state
        self.num_wordcloud_words = num_wordcloud_words

    def create_model(self, token_lists, output_path):
        """ Generate model & evaluation file to a given output path

        :param token_lists: list of document tokens
        :param output_path: path to save model, dictionary, corpus, evaluation files
        """
        try:

            # create topic model & fit token_lists
            tm = TM(
                method=self.method, random_state=self.random_state, pre_trained_name=self.pre_trained_name,
                n_neighbors=self.n_neighbors, n_components=self.n_components, min_dist=self.min_dist,
                umap_metric=self.umap_metric, min_cluster_size=self.min_cluster_size, cluster_metric=self.cluster_metric,
                cluster_selection_method=self.cluster_selection_method, prediction_data=self.prediction_data
            )

            tm.fit(token_lists)

            # define output path
            subdir = "{}_k_{}_{}_{}".format(self.lang_code, str(tm.k), self.method, self.version)
            models_path = output_path + "/models/" + subdir
            dict_path = output_path + "/dictionary/"
            corpus_path = output_path + "/corpus/"
            eval_path = output_path + "/evaluation/" + subdir

            # create directories
            make_dirs(output_path)
            make_dirs(models_path)
            make_dirs(dict_path)
            make_dirs(corpus_path)
            make_dirs(eval_path)
            make_dirs(eval_path + "/wordcloud")

            # extract topics
            topics = self.extract_topics(tm.tf_idf, tm.feature_names, tm.labels)
            topic_words = self.get_topic_words(topics)

            # evaluate topics
            c_v = self.compute_coherence(topic_words, tm.dictionary, tm.corpus, token_lists, measure='c_v')
            u_mass = self.compute_coherence(topic_words, tm.dictionary, tm.corpus, token_lists, measure='u_mass')
            silhouette = self.compute_silhouette(tm.umap_embeddings, tm.labels)

            # hyperparameters & evaluation metrics
            metrics = {
                "k"                         : tm.k,
                "lang_code"                 : self.lang_code,
                "num_docs"                  : len(token_lists),
                "method"                    : self.method,
                "version"                   : self.version,
                "norm_type"                 : self.norm_type,
                "pre_trained_name"          : self.pre_trained_name,
                "random_state"              : self.random_state,
                "n_neighbors"               : self.n_neighbors,
                "n_components"              : self.n_components,
                "min_dist"                  : self.min_dist,
                "umap_metric"               : self.umap_metric,
                "min_cluster_size"          : self.min_cluster_size,
                "cluster_metric"            : self.cluster_metric,
                "cluster_selection_method"  : self.cluster_selection_method,
                "c_v"                       : c_v,
                "u_mass"                    : u_mass,
                "silhouette"                : silhouette
            }

            # save topic model, dictionary, corpus
            tm.dictionary.save(dict_path + "/dict.gensim")
            dump_to_pickle(tm.corpus, corpus_path + "/corpus.pkl")
            dump_to_pickle(tm.sentences, eval_path + "/sentences.pkl")
            dump_to_pickle(tm.embeddings, eval_path + "/embeddings.pkl")
            dump_to_pickle(topics, models_path + "/topics.pkl")
            dump_to_pickle(tm.cluster_model, models_path + "/cluster.pkl")
            dump_to_pickle(tm.umap_model, models_path + "/umap.pkl")

            # save topic terms
            self.save_topic_terms(topics, eval_path + '/topic_terms.txt')

            # save metrics, topic_terms dataframe
            dump_to_json(metrics, eval_path + "/evaluation.json", sort_keys=False)

            # save topic visualization, topic wordclouds
            save_topic_visualization(tm.embeddings, tm.labels, eval_path + "/topics.png")
            save_topic_wordclouds(topics, self.num_wordcloud_words, eval_path + "/wordcloud")

        except Exception:
            logging.error('error occured', exc_info=True)

    def compute_coherence(self, topics, dictionary, corpus, token_lists, measure):
        """ Compute coherence score for a given topic model

        :param topics: topics
        :param dictionary: generated dictionary
        :param corpus: generated bow corpus
        :param token_lists: list of document tokens
        :param measure: coherence score
        """
        coherence = 0.0
        cm = CoherenceModel(
            topics=topics,
            dictionary=dictionary,
            corpus=corpus,
            texts=token_lists,
            coherence=measure
        )
        coherence = cm.get_coherence()
        coherence = "{0:.3f}".format(coherence)
        return coherence

    def compute_silhouette(self, umap_embeddings, labels):
        """ Compute silhouette score for cluster distance and relevance

        :param umap_embeddings:
        """
        silhouette = silhouette_score(umap_embeddings, labels)
        silhouette = float("{0:.3f}".format(silhouette))
        return silhouette

    def extract_topics(self, tf_idf, feature_names, labels):
        """ Extract topic words and their c-TF-IDF score within each cluster

        :param tf_idf: tf_idf scores
        :param feature_names: dictionary words
        :param labels: cluster labels
        :return: dictionary (topic: top_n_words)
        Sample:
            {
                114 : [
                        ("cough", 0.20982713670636632)
                        ("fever", 0.0538832539146329)
                        ("throat", 0.02416355064071378)
                        ("case", 0.023852171240297093)
                        ("fever_cough", 0.02382326255347972)
                        ("illness", 0.023609866349192463)
                        ("radiograph", 0.022765546069615183)
                        ("syndrome", 0.02179170899590873)
                        ("pneumonia", 0.02169239788185851)
                ]
            }
        """
        labels = sorted(list(set(labels)))
        tf_idf_transposed = tf_idf.T
        indices = tf_idf_transposed.argsort()[::-1]

        topics = dict()
        for i, label in enumerate(labels):
            top_n_words = []
            for j in indices[i]:
                word = feature_names[j]
                c_tf_idf = tf_idf_transposed[i][j]
                top_n_words.append((word, c_tf_idf))
            top_n_words = top_n_words[::-1]
            topics[label] = top_n_words
        return topics

    def get_topic_words(self, topics):
        """ Get list of topic words from topics dictionary

        :param topics: dictionary (topic:top_n_words)
        :return: list of list of topic words
        Sample:
            [
                [
                    "risk", "increase", "treatment", "case", "system"
                ],
                [
                    "obesity", "weight", "surgery", "complication", "procedure"
                ],
                [
                    "tourism", "destination", "arrival", "paper", "industry"
                ]

            ]
        """
        topic_words = []
        for topic, top_n_words in topics.items():
            words = [word for word, c_tf_idf in top_n_words]
            topic_words.append(words)
        return topic_words

    def get_topic_terms_df(self, topics):
        """ Generate topic_term dataframe for a given model

        :param topics: dictionary (topic:top_n_words)
        :return: dataframe
        """
        labels = list(topics.keys())

        topic_terms = []
        for topic, top_n_words in topics.items():
            top_n_words = sorted(top_n_words, key=lambda x: x[1], reverse=True)[:self.num_words]
            terms = [term for term, c_tf_idf in top_n_words]
            terms = ", ".join(terms)
            topic_terms.append(terms)

        topic_terms_df = pd.DataFrame()
        topic_terms_df['id'] = labels
        topic_terms_df['Topic terms'] = topic_terms
        return topic_terms_df

    def save_topic_terms(self, topics, output_path):
        """ Generate and save topic_terms dataframe to a given output_path

        :param topics: dictionary (topic:top_n_words)
        :param output_path: output path
        """
        topic_terms_df = self.get_topic_terms_df(topics)
        topic_terms_df.to_string(output_path, index=False)

if __name__ == "__main__":

    chunks_prep_path = abs_path = os.path.dirname(os.path.abspath("__file__")) + "/chunks_prep"
    output_path = os.path.dirname(os.path.abspath("__file__")) + "/output"

    token_lists = []
    for filename in os.listdir(chunks_prep_path):
        chunk_no = filename.split("_")[1]
        token_list = load_from_pickle(chunks_prep_path + "/" + filename + "/token_lists.pkl")
        token_lists.append(token_list)

    token_lists = [item for sublist in token_lists for item in sublist]
    print("NUM DOCS: ", len(token_lists))

    start_time = time.time()
    ev = Evaluation(lang_code="en", method="BERT", version="1.1", num_words=15)
    ev.create_model(token_lists, output_path=output_path)
    print("--- BERT: %s seconds ---" % (time.time() - start_time))


"""
Buffer that also stores advantage estimates as described in PPO.

Author: Ian Char
Date: 4/7/2021
"""
import numpy as np
import torch
import pudb
import copy
from rlkit.data_management.onpolicy_buffer import OnPolicyBuffer
from rlkit.torch.core import torch_ify

class AdvantageReplayBuffer(OnPolicyBuffer):
    """Buffer that also stores advantages."""

    def __init__(
        self,
        max_replay_buffer_size,
        env,
        value_function,
        intrinsic_value_function,
        tdlambda=0.95,
        discount=0.99,
        intrinsic_discount = 0.9999,
        target_lookahead=10,
        use_dones_for_rnd_critic=False,
    ):
        self._max_replay_buffer_size = max_replay_buffer_size
        self._observation_dim = env.observation_space.low.size
        self._action_dim = env.action_space.low.size
        self._value_function = value_function
        self._intrinsic_value_function = intrinsic_value_function
        self._tdlambda = tdlambda
        self._discount = discount
        self._intrinsic_discount = intrinsic_discount
        self.use_dones_for_rnd_critic = use_dones_for_rnd_critic
        self._target_lookahead = target_lookahead
        self.clear_buffer()

    def clear_buffer(self):
        """Clear the replay buffers."""
        # Set up buffers.
        self._observations = np.zeros((self._max_replay_buffer_size,
                                       self._observation_dim))
        self._next_obs = np.zeros((self._max_replay_buffer_size,
                                   self._observation_dim))
        self._actions = np.zeros((self._max_replay_buffer_size,
                                  self._action_dim))
        self._rewards = np.zeros((self._max_replay_buffer_size, 1))
        self._intrinsic_rewards = np.zeros((self._max_replay_buffer_size, 1))
        self._intrinsic_rewards_raw = np.zeros((self._max_replay_buffer_size, 1))
        self._advantages = np.zeros((self._max_replay_buffer_size, 1))
        self._intrinsic_advantages = np.zeros((self._max_replay_buffer_size, 1))
        self._targets = np.zeros((self._max_replay_buffer_size, 1))
        self._intrinsic_targets = np.zeros((self._max_replay_buffer_size, 1))
        self._logpis = np.zeros((self._max_replay_buffer_size, 1))
        # self._terminals[i] = a terminal was received at time i
        self._terminals = np.zeros((self._max_replay_buffer_size, 1),
                                   dtype='uint8')
        self._top = 0
        self._size = 0
        self.running_std_intrinsic_return = 0
        self.running_var_intrinsic_return = 0
        self.running_mean_intrinsic_return = 0
        self.num_return_samples = 0

    def random_batch(self, batch_size):
        indices = np.random.randint(0, self._size, batch_size)
        batch = dict(
            observations=self._observations[indices],
            actions=self._actions[indices],
            rewards=self._rewards[indices],
            intrinsic_rewards=self._intrinsic_rewards[indices],
            intrinsic_rewards_raw=self._intrinsic_rewards_raw[indices],
            terminals=self._terminals[indices],
            next_observations=self._next_obs[indices],
            advantages=self._advantages[indices],
            intrinsic_advantages=self._intrinsic_advantages[indices],
            targets=self._targets[indices],
            intrinsic_targets=self._intrinsic_targets[indices],
            logpis=self._logpis[indices],
        )
        return batch

    def add_path(self, path):
        obs = path['observations']
        acts = path['actions']
        rews = path['rewards']
        intrinsic_rews = path['intrinsic_rewards']
        intrinsic_rews_raw = copy.deepcopy(intrinsic_rews)
        nxts = path['next_observations']
        dones = path['terminals']
        pis = np.array([ai['logpi'] for ai in path['agent_infos']])
        path_len = len(obs)

        # Normalise the intrinsic rewards
        # NOTE: this is incorrect, it should be normalizing by a running std. dev. of the intrinsic returns, not the rewards
        self.num_return_samples += 1
        self.running_mean_intrinsic_return += (1/self.num_return_samples) * (np.sum(intrinsic_rews) - self.running_mean_intrinsic_return)

        if self.num_return_samples >= 2:
            self.running_var_intrinsic_return = ((self.num_return_samples - 2)/(self.num_return_samples - 1)) * self.running_var_intrinsic_return
            self.running_var_intrinsic_return += (1/self.num_return_samples) * ((np.sum(intrinsic_rews) - self.running_mean_intrinsic_return)**2)
            self.running_std_intrinsic_return = np.sqrt(self.running_var_intrinsic_return)

            intrinsic_rews /= (self.running_std_intrinsic_return + 1e-8)

        # Compute the advantages.
        with torch.no_grad():
            vals = self._value_function(torch_ify(obs)).cpu().numpy()
            nxt_vals = self._value_function(torch_ify(nxts)).cpu().numpy()
        deltas = rews + self._discount * nxt_vals * (1 - dones) - vals
        advantages = np.zeros(path_len)
        for idx, delta in enumerate(deltas.flatten()[::-1]):
            if idx == 0:
                advantages[-1] = delta
            else:
                advantages[path_len - 1 - idx] = (delta
                        + self._discount * self._tdlambda
                        * advantages[path_len - idx])
        targets = advantages + vals.flatten()

        # Compute the intrinsic advantages.
        with torch.no_grad():
            intrinsic_vals = self._intrinsic_value_function(
                torch_ify(obs)).cpu().numpy()
            nxt_intrinsic_vals = self._intrinsic_value_function(
                torch_ify(nxts)).cpu().numpy()
        intrinsic_deltas = intrinsic_rews + self._intrinsic_discount * \
                nxt_intrinsic_vals * (1 - (dones * self.use_dones_for_rnd_critic)) - intrinsic_vals
        intrinsic_advantages = np.zeros(path_len)
        for idx, intrinsic_delta in enumerate(intrinsic_deltas.flatten()[::-1]):
            if idx == 0:
                intrinsic_advantages[-1] = intrinsic_delta
            else:
                intrinsic_advantages[path_len - 1 - idx] = (intrinsic_delta
                        + self._intrinsic_discount * self._tdlambda
                        * intrinsic_advantages[path_len - idx])
        intrinsic_targets = intrinsic_advantages + intrinsic_vals.flatten()

        for o, a, r, ir, irr, n, d, ad, iad, pi, t , it in zip(obs, acts, rews,
                                            intrinsic_rews, intrinsic_rews_raw, nxts, dones,
                                            advantages, intrinsic_advantages,
                                            pis, targets, intrinsic_targets):
            self._observations[self._top] = o
            self._actions[self._top] = a
            self._rewards[self._top] = r
            self._intrinsic_rewards[self._top] = ir
            self._intrinsic_rewards_raw[self._top] = irr
            self._terminals[self._top] = d
            self._next_obs[self._top] = n
            self._advantages[self._top] = ad
            self._intrinsic_advantages[self._top] = iad
            self._targets[self._top] = t
            self._intrinsic_targets[self._top] = it
            self._logpis[self._top] = pi

            self._advance()

    def add_paths(self, paths):
        for path in paths:
            self.add_path(path)

    def _advance(self):
        self._top = (self._top + 1) % self._max_replay_buffer_size
        if self._size < self._max_replay_buffer_size:
            self._size += 1


import skimage.io as io
import skimage.transform as skt
import numpy as np
from PIL import Image
from src.models.class_patcher import patcher
from src.utils.imgproc import *

class patcher(patcher):
    def __init__(self, body='./body/body_noy.png', **options):
        super().__init__('ノイ', body=body, pantie_position=[147, 133], **options)
        self.mask = io.imread('./mask/mask_noy.png')

    def convert(self, image):
        pantie = np.array(image)

        # prepare for moving from hip to front
        patch = np.copy(pantie[-140:-5, 546:, :])
        patch = skt.resize(patch[::-1, ::-1, :], (240, 60), anti_aliasing=True, mode='reflect')
        [pr, pc, d] = patch.shape

        # Affine transform matrix
        pantie = np.pad(pantie, [(0, 0), (100, 0), (0, 0)], mode='constant')
        arrx = np.zeros(100)
        arry = np.zeros(100)
        arry[:30] += np.linspace(70,0,30)
        arry[30:] += np.linspace(0,200,70)
        arry[65:] -= np.linspace(0,170,35)
        pantie = affine_transform_by_arr(pantie, arrx, arry)[:,53:-65]
        pantie[147:147 + pr, :pc, :] = patch
        pantie = np.bitwise_and(np.uint8(pantie*255), self.mask)
        pantie = skt.resize(pantie, (int(pantie.shape[0]*3.65), int(pantie.shape[1]*3.13)), anti_aliasing=True, mode='reflect')[:,8:]
        pantie = np.concatenate((pantie[:,::-1],pantie),axis=1)

        return Image.fromarray(np.uint8(pantie*255))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Sentiment Analysis on Binary sentiments of a Restaurant : 0 for Bad and 1 for Good

@author: kaustabh

License : MIT (c) Kaustabh Ganguly 2018

"""

# Importing the libraries
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# Importing the dataset
dataset = pd.read_csv('Restaurant_Reviews.tsv', delimiter = '\t', quoting = 3)

# Cleaning the texts
import re
import nltk
#nltk.download('stopwords')
from nltk.corpus import stopwords
from nltk.stem.porter import PorterStemmer
corpus = []
for i in range(0, 1000):
    review = re.sub('[^a-zA-Z]', ' ', dataset['Review'][i])
    review = review.lower()
    review = review.split()
    ps = PorterStemmer()
    review = [ps.stem(word) for word in review if not word in set(stopwords.words('english'))]
    review = ' '.join(review)
    corpus.append(review)

# Creating the Bag of Words model
from sklearn.feature_extraction.text import CountVectorizer
cv = CountVectorizer(max_features = 600)
X = cv.fit_transform(corpus).toarray()
y = dataset.iloc[:, 1].values

# Splitting the dataset into the Training set and Test set
from sklearn.cross_validation import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state = 0)

# Fitting Naive Bayes to the Training set
from sklearn.naive_bayes import BernoulliNB
classifier = BernoulliNB()
classifier.fit(X_train, y_train)

# Predicting the Test set results
y_pred = classifier.predict(X_test)

# Making the Confusion Matrix
from sklearn.metrics import confusion_matrix
cm = confusion_matrix(y_test, y_pred)


"""
Usage: motif_extractor <json_file>
"""
# @Author  : Yu Li & Siyuan Chen
# @Software: PyCharm

import scipy.signal as signal
from lib.classification_models import *
from tqdm import tqdm
from collections import OrderedDict
import deeplift.conversion.kerasapi_conversion as kc
from deeplift.layers import NonlinearMxtsMode
from deeplift.visualization.viz_sequence import *
import os,docopt
import json

root = os.path.dirname(os.path.abspath(__file__))
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

def one_hotidx_2_base(idx):
    if idx == 0 :
        return "A"
    elif idx == 1 :
        return "C"
    elif idx == 2 :
        return "G"
    elif idx == 3 :
        return "T"

def one_hot2_seq(one_hot_array):
    seq = ""
    _,non_zero_idx = np.nonzero(one_hot_array)
    seq = seq.join([one_hotidx_2_base(base) for base in non_zero_idx])
    return seq

default_colors = {0:'green', 1:'blue', 2:'orange', 3:'red'}
default_plot_funcs = {0:plot_a, 1:plot_c, 2:plot_g, 3:plot_t}
def save_plot_weights(array,
                 figsize=(20,2),
                 height_padding_factor=0.2,
                 length_padding=1.0,
                 subticks_frequency=1.0,
                 colors=default_colors,
                 plot_funcs=default_plot_funcs,
                 highlight={},
                 enriched_factor = 0,
                 picname = None):
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111)
    plot_weights_given_ax(ax=ax, array=array,
        height_padding_factor=height_padding_factor,
        length_padding=length_padding,
        subticks_frequency=subticks_frequency,
        colors=colors,
        plot_funcs=plot_funcs,
        highlight=highlight)
    plt.title('Enriched Factor :{}'.format(enriched_factor), fontsize='large', fontweight = 'bold')
    plt.savefig(picname)

def get_rate(args,rate_mat_file):
    rate_folder_name = os.path.join(root, 'dataset',
                                    args["dataset"],
                                    str(args["input_length"])
                                    )
    data = json.load(open(os.path.join(rate_folder_name,"hot_data_list.json")))
    rate_mat = []

    if len(args["dataset"]) == 3:
        char_rate = "Rate"
    elif args["dataset"] == "mouse_cell_2016":
        char_rate = "BPM"
    else:
        char_rate = "rate"

    if args["dataset"] in ["nature_genetics_2008","yeast_2008" ]:

        for i in range(len(data)):
            rate_mat.append(10)
    else:
        for d in data:
            rate_mat.append(d[char_rate])

    np.save(rate_mat_file, np.array(rate_mat))
    return np.array(rate_mat)

def get_score(args,score_mat_file):
    data_folder_name = os.path.join(root, 'dataset',
                                    args["dataset"],
                                    str(args["input_length"])
                                    )

    data = np.load(os.path.join(data_folder_name,"X.npy"))
    print("=> {} loaded ".format(os.path.join(root,data_folder_name,"X.npy")))
    Y = np.load(os.path.join(data_folder_name,"Y.npy"))

    num_hot = np.count_nonzero(Y, axis=None)

    hot_data_ont_hot = np.array(data[0:num_hot])

    model_file_name = os.path.join(root, 'model',
                                   args["dataset"],
                                   args["model"],
                                   str(args["input_length"]) + "_" + str(args["input_height"]),
                                   args["output_prefix"]
                                   )

    keras_model_weights = model_file_name + ".h5"
    keras_model_yaml = model_file_name + ".yaml"
    keras_model_json = model_file_name + ".json"

    model = CNN(epochs=args["epochs"], args=args, loss="binary_crossentropy")
    model.load_weights(keras_model_weights)

    # print("=> {} loaded ".format(keras_model_json))
    # deeplift_model = kc.convert_model_from_saved_files(
    #     h5_file=keras_model_weights,
    #     json_file=keras_model_json)

    guided_backprop = kc.convert_model_from_saved_files(
        h5_file=keras_model_weights,
        json_file=keras_model_json,
        nonlinear_mxts_mode=NonlinearMxtsMode.GuidedBackprop)
    guided_backprop_scoring_func = guided_backprop.get_target_contribs_func(find_scores_layer_idx=0,
                                                                            target_layer_idx=-2)

    method_to_task_to_scores = OrderedDict()
    method_name = "guided_backprop"
    score_func = guided_backprop_scoring_func
    print("on method", method_name)
    method_to_task_to_scores[method_name] = OrderedDict()
    background = OrderedDict([('A', 0.3), ('C', 0.2), ('G', 0.2), ('T', 0.3)])

    for task_idx in [0,1]:
        scores = np.array(score_func(
            task_idx=task_idx,
            input_data_list=[hot_data_ont_hot],
            input_references_list=[
                np.array([background['A'],
                          background['C'],
                          background['G'],
                          background['T']])[None, None, :]],
            batch_size=200,
            progress_update=None))

    for task_idx in [0,1]:
        scores = np.array(score_func(
            task_idx=task_idx,
            input_data_list=[hot_data_ont_hot],
            input_references_list=[
                np.array([background['A'],
                          background['C'],
                          background['G'],
                          background['T']])[None, None, :]],
            batch_size=200,
            progress_update=None))
        assert scores.shape[2] == 4
        scores = np.sum(scores, axis=2)
        method_to_task_to_scores[method_name][task_idx] = scores

    score_mat = []
    task_idx = [(0, seq) for seq in range(num_hot)]

    for task, idx in tqdm(task_idx):

        method_name = 'guided_backprop'
        scores = method_to_task_to_scores[method_name][task]

        scores_for_idx = scores[idx]
        original_onehot = hot_data_ont_hot[idx]
        scores_for_idx = original_onehot * scores_for_idx[:, None]
        score_mat.append(scores_for_idx)
    np.save(score_mat_file, np.array(score_mat))
    return np.array(score_mat)

def Sliding_Window(scores_for_idx, k):

    top_score = 0
    # Sliding window
    motif = {
        "k_mer": np.zeros((k,4)),  # matrix
        "sequence": "",  # string
        "length": k,
        "start": 0,
        "end": 0,
        "score": 0,
        "count": 0
    }
    for i in range(0, 1000 - k):
        score = np.average(scores_for_idx[i:i + k])
        if score > top_score:
            top_score = score
            motif["k_mer"] = scores_for_idx[i:i + k]
            motif["score"] = top_score
            motif["start"] = int(i)
            motif["end"] = int(i + k)
            motif["sequence"] = one_hot2_seq(scores_for_idx[i:i + k])

    return motif

def Low_pass_filter_extractor(scores_for_idx, recomb_rate, filter_factor):
    # inpt.shape [1000,4]

    motif_list = []
    sample_non_zero_idx, sample_non_zero_ATCG = scores_for_idx.nonzero()
    sample_signal = np.zeros((1000,))

    for i in range(len(sample_non_zero_idx)):
        pos = sample_non_zero_idx[i]
        sample_signal[pos] = scores_for_idx[pos][sample_non_zero_ATCG[i]]

    b, a = signal.butter(8, filter_factor, 'lowpass')
    filtered_signal = signal.filtfilt(b, a, sample_signal)
    # inpt.shape [1000]

    peaks, _ = signal.find_peaks(filtered_signal, prominence=0.06)
    valleys, _ = signal.find_peaks(filtered_signal * (-1), width=1)
    start_idx_list = []
    end_idx_list = []
    for peak in peaks:
        reverse_peak = np.array([valleys[v] - peak for v in range(len(valleys))])
        start = valleys[np.where(reverse_peak < 0, reverse_peak, -np.inf).argmax()]
        end = valleys[np.where(reverse_peak > 0, reverse_peak, np.inf).argmin()]
        start_idx_list.append(start)
        end_idx_list.append(end)

    for i in range(len(start_idx_list)):
        motif = {
            "k_mer":scores_for_idx[start_idx_list[i]:end_idx_list[i]],  # matrix
            "sequence": "",  # string
            "recomb_rate":recomb_rate,
            "length": 0,
            "start": 0,
            "end": 0,
            "score": 0,
            "enriched_factor":0,
            "count": 0,
        }

        motif["start"] = start_idx_list[i]
        motif["end"] = end_idx_list[i]
        if motif["end"] < motif["start"]:
            print("Wrong peak, Continue")
            continue
        motif["score"] = np.average(sample_signal[motif["start"]:motif["end"]])
        motif["enriched_factor"] = motif["score"]/np.average(sample_signal)
        motif["length"] = int(end_idx_list[i] - start_idx_list[i])
        motif["sequence"] = one_hot2_seq(scores_for_idx[start_idx_list[i]:end_idx_list[i]])
        motif_list.append(motif)

    return motif_list

def count_alg(motif, sequence_list):
    count = 0
    for sequence in sequence_list:
        count += sequence.count(motif)

    return count

class NumpyEncoder(json.JSONEncoder):
    """ Special json encoder for numpy types """
    def default(self, obj):
        if isinstance(obj, (np.int_, np.intc, np.intp, np.int8,
                            np.int16, np.int32, np.int64, np.uint8,
                            np.uint16, np.uint32, np.uint64)):
            return int(obj)
        elif isinstance(obj, (np.float_, np.float16, np.float32,
                              np.float64)):
            return float(obj)
        elif isinstance(obj, (np.ndarray,)):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)
if __name__ == '__main__':
    args = docopt.docopt(__doc__)
    args = json.load(open(args["<json_file>"]))

    # Do counting
    target_length = 13
    result_root = os.path.join(root,
                               "motifs",
                               args["dataset"]
                               )
    if not os.path.exists(result_root):
        os.makedirs(result_root)

    # Get scores
    score_mat_file = os.path.join(result_root,"scores.npy")
    rate_mat_file = os.path.join(result_root,"recomb_rate.npy")

    if not os.path.exists(rate_mat_file):
        recomb_rates = get_rate(args,rate_mat_file)
    else:
        recomb_rates = np.load(rate_mat_file)

    if not os.path.exists(score_mat_file):
        scores = get_score(args,score_mat_file)
    else:
        scores = np.load(score_mat_file)

    data_folder_name = os.path.join(root, 'dataset',
                                    args["dataset"],
                                    str(args["input_length"])
                                    )

    labels = np.load(os.path.join(data_folder_name,"Y.npy"))

    num_data = np.count_nonzero(labels, axis=None)

    task_idx = [(0, seq) for seq in range(num_data)]
    # Extract Motifs

    # filter_list = [0.1]
    if "_2014" in args["dataset"]:
        filter_list = [0.4, 0.2, 0.1]
        num_save = 20
    else:
        num_save = 500
        filter_list = [0.4, 0.2, 0.1]

    for filter_v in filter_list:
        motif_list = []
        for task, idx in tqdm(task_idx):
            scores_for_idx = scores[idx]
            recomb_rate = recomb_rates[idx]
            # scores_for_idx.shape =  [1000,4]
            # motif_list.append(Sliding_Window(scores_for_idx, target_length ))
            m_list = Low_pass_filter_extractor(scores_for_idx,
                                               recomb_rate,
                                               filter_factor = filter_v)
            for m in m_list:
                motif_list.append(m)

        motif_list = sorted(motif_list, key=lambda d: d["enriched_factor"], reverse=True)

        # Draw motif
        num_draw = 20
        for i in range(0, num_draw):
            pic_root = os.path.join(result_root,"filter{}".format(filter_v))
            if not os.path.exists(pic_root):
                os.mkdir(pic_root)
            print(motif_list[i]["sequence"])
            start = int(motif_list[i]["start"])
            end = int(motif_list[i]["end"])
            print("Whole Sequence ranking {}:".format(i + 1))
            #     highlight_seq = {'blue': [
            #         (start, end)]
            #     }
            save_plot_weights(motif_list[i]["k_mer"],
                              subticks_frequency=1,
                              enriched_factor=motif_list[i]["enriched_factor"],
                              picname= os.path.join(pic_root,'filter{}_rank_{}.jpg'.format(filter_v,i + 1)))
            print('=> {} saved'.format(os.path.join(pic_root,'filter{}_rank_{}.jpg'.format(filter_v,i + 1))))

        saved_file = os.path.join(result_root,'filter{}_motifs.json'.format(filter_v))
        save_list = []
        for i in range(0,num_save):
            motif_list[i]["rank"] = i+1
            save_list.append(motif_list[i])
        print('saving ==>{}'.format(saved_file))
        with open(saved_file, 'w') as fw:
            js.dump(save_list, fw,cls=NumpyEncoder, indent=4)

    #     save_plot_weights(sequence_list[i]["sequence"],
    #                       subticks_frequency=100,
    #                       highlight=highlight_seq,
    #                       confidence_score = sequence_list[i]["score"],
    #
    #                       picname = os.path.join(result_root,'sequence_rank_{}.jpg'.format(i + 1)))
    #
    #