# -*- coding:UTF8 -*-

"""Image bounding-box extraction process."""

import os
import sys

import cv2
import numpy as np

import torch
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer, default_setup, launch

from utils.extract_utils import get_image_blob
from models import add_config
from models.bua.box_regression import BUABoxes
from models.bua import add_bottom_up_attention_config
from detectron2.layers.nms import nms

class BboxExtractor:
    def __init__(self, cfg_file, gpu_id = 0):

        self.cfg_file = cfg_file
        self.gpu_id = gpu_id
        self.cfg = get_cfg()
        add_bottom_up_attention_config(self.cfg, True)
        self.cfg.merge_from_file(self.cfg_file)
        self.cfg.freeze()
        default_setup(self.cfg, None)

        self.bbox_extract_model = DefaultTrainer.build_model(self.cfg)
#        self.bbox_extract_model.cuda(gpu_id)
        bbox_extract_model_dict = self.bbox_extract_model.state_dict()
        bbox_extract_checkpoint_dict = torch.load(self.cfg.MODEL.WEIGHTS, map_location=torch.device('cuda:0'))['model']
        bbox_extract_checkpoint_dict = {k:v for k, v in bbox_extract_checkpoint_dict.items() if k in bbox_extract_model_dict}
        bbox_extract_model_dict.update(bbox_extract_checkpoint_dict)
        self.bbox_extract_model.load_state_dict(bbox_extract_model_dict)
        # self.bbox_extract_model = torch.nn.DataParallel(self.bbox_extract_model, device_ids=self.gpus)
        self.bbox_extract_model.eval()

    def clean_bbox(self, dataset_dict, boxes, scores):
        MIN_BOXES = self.cfg.MODEL.BUA.EXTRACTOR.MIN_BOXES
        MAX_BOXES = self.cfg.MODEL.BUA.EXTRACTOR.MAX_BOXES
        CONF_THRESH = self.cfg.MODEL.BUA.EXTRACTOR.CONF_THRESH

        scores = scores[0]
        boxes = boxes[0]
        num_classes = scores.shape[1]
        boxes = BUABoxes(boxes.reshape(-1, 4))
        boxes.clip((dataset_dict['image'].shape[1]/dataset_dict['im_scale'], dataset_dict['image'].shape[2]/dataset_dict['im_scale']))
        boxes = boxes.tensor.view(-1, num_classes*4)  # R x C x 4

        cls_boxes = torch.zeros((boxes.shape[0], 4))
        for idx in range(boxes.shape[0]):
            cls_idx = torch.argmax(scores[idx, 1:]) + 1
            cls_boxes[idx, :] = boxes[idx, cls_idx * 4:(cls_idx + 1) * 4]

        max_conf = torch.zeros((scores.shape[0])).to(scores.device)
        for cls_ind in range(1, num_classes):
                cls_scores = scores[:, cls_ind]
                keep = nms(cls_boxes, cls_scores, 0.3)
                max_conf[keep] = torch.where(cls_scores[keep] > max_conf[keep],
                                                 cls_scores[keep],
                                                 max_conf[keep])

        keep_boxes = torch.argsort(max_conf, descending=True)[:MAX_BOXES]
        image_bboxes = cls_boxes[keep_boxes]

        return image_bboxes

    def extract_bboxes(self, img_path):
        im = cv2.imread(img_path)
        if im is None:
            print("img is None!")

            return 501, None
        else:
            dataset_dict = get_image_blob(im, self.cfg.MODEL.PIXEL_MEAN)
#            dataset_dict["image"] = dataset_dict["image"].cuda(self.gpu_id)
#            print(dataset_dict["image"])
#            print("-----model-----")
#            print(self.bbox_extract_model)
            with torch.set_grad_enabled(False):
                boxes, scores = self.bbox_extract_model([dataset_dict])
            boxes = [box.cpu() for box in boxes]
            scores = [score.cpu() for score in scores]
            boxes = self.clean_bbox(dataset_dict, boxes, scores)

            return 200, boxes # boxes type tensor

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_path', type=str, dest="img_path", default="test_data/test.png")
    parser.add_argument('--out_path', type=str, dest="output_file", default="test_data/test.npz")
    args = parser.parse_args()
    bbx_extr = BboxExtractor('configs/bua-caffe/extract-bua-caffe-r101.yaml')
    code, bboxes = bbx_extr.extract_bboxes(args.img_path)
    np.savez_compressed(args.output_file, bbox=bboxes)
    print(code)
    np_bbox = bboxes.numpy().astype(np.int32)
    print(np_bbox)
    print(bboxes.shape)


import streamlit as st
import yaml
import loader
import numpy as np
import pandas as pd

# import sys
from models import IndicatorType, IndicatorCards, ProductCards, DimensionCards

# import pdf_report.pdfgen as pdfgen
import pages.simulacovid as sm
import pages.saude_em_ordem as so
import pages.distancing as ds
import pages.onda_covid as oc
import pages.estudo as estudo
import plots
import utils
import amplitude
import session

from streamlit.server.server import Server
import os

import bisect
import math
from bokeh.models.widgets import Div

# Arruma os indicadores do filtro
# Fix indicators filters
def fix_type(x, group, position):
    if type(x) == np.ndarray:
        if "- " in x:
            return "- "
        else:
            return " a ".join(
                [str(round(i, 1)) if type(i) != str else i.strip() for i in x]
            )

    if x == "- ":
        return x

    if position == "last_updated":
        return pd.to_datetime(str(x)).strftime("%d/%m/%Y")

    if group == "situation" and position == "display":
        return round(float(x), 2)

    if group == "trust" and position == "display":
        return int(round(10 * (1 - x), 0))  # muda para não notificado

    if group == "situation" and position == "left_display":
        return str(int(x)) + " dias"

    if group == "capacity" and position == "display":
        # TODO -> VOLTAR PARA PROJECAO DE LEITOS
        # return utils.dday_preffix(x)
        return int(x)

    if (type(x) == str) or (type(x) == np.int64) or (type(x) == int):
        return x

    if type(x) == np.float64:
        return int(x)

# Atualiza os indicadores para arrumar valores e disposicao
# Updates the indicators to adjust values and disposition
def update_indicators(indicators, data, config, user_input, session_state):
    # TODO: indicadores quando cidade não posssui dados
    for group in config["br"]["indicators"].keys():
        # Arrumar valores para cada posição
        # Fix values for each position
        dic_indicators = dict()
        if group == 'vacina':
            if "state" in user_input["place_type"]:
                dado = pd.read_csv('http://datasource.coronacidades.org/br/states/vacina')
                dado = dado[dado["state_name"] == user_input["state_name"]]
            elif "city" in user_input["place_type"]:
                dado = pd.read_csv('http://datasource.coronacidades.org/br/cities/vacina')
                dado = dado[dado["city_id"] == user_input["city_id"]]
            else:
                dado = pd.read_csv('http://datasource.coronacidades.org/br/health_region/vacina')
                dado = dado[dado['health_region_id'] == user_input["health_region_id"]]
            indicators[group].last_updated = dado.last_updated.fillna("-").values[0]
            indicators[group].imunizados = dado.imunizados.fillna("-").values[0]
            indicators[group].nao_vacinados = dado.nao_vacinados.fillna("-").values[0]
            indicators[group].perc_imunizados = dado.perc_imunizados.fillna("-").values[0]
            indicators[group].perc_vacinados = dado.perc_vacinados.fillna("-").values[0]
        else:
            for position in config["br"]["indicators"][group].keys():
                # Indicadores de Filtros
                # Filter indicators
                if config["br"]["indicators"][group][position] != "None":
                    dic_indicators[position] = fix_type(
                        data[config["br"]["indicators"][group][position]]
                        .fillna("- ")
                        .values[0],
                        group,
                        position,
                    )
                else:
                    dic_indicators[position] = "None"
            if np.isnan(data[config["br"]["indicators"][group]["risk"]].values[0]):
                dic_indicators["risk"] = "nan"
                if group == "rt":  # doesn't show values
                    indicators[group].right_display = "- "
                    indicators[group].left_display = "- "

            # TODO: melhroar aqui para não ter que passar os dados de um dic para outro
            indicators[group].risk = dic_indicators["risk"]
            indicators[group].left_display = dic_indicators["left_display"]
            indicators[group].right_display = dic_indicators["right_display"]
            indicators[group].display = dic_indicators["display"]
            indicators[group].last_updated = dic_indicators["last_updated"]

    # Caso o usuário altere os casos confirmados, usamos esse novo valor para a estimação
    # TODO: vamos acabar com o user_iput e manter só session_state?
    if (session_state.number_cases is not None) and (
        session_state.number_cases != user_input["population_params"]["I_compare"]
    ):
        user_input["population_params"]["I"] = session_state.number_cases
        user_input["population_params"]["D"] = session_state.number_cases

    return indicators

# Atualiza as escolhas com o ID da selecao do usúario
# Update choices with user selection ID
def update_user_input_ids(data, user_input):

    user_input["state_num_id"] = data["state_num_id"].values[0]
    user_input["health_region_id"] = False
    user_input["city_id"] = False

    if user_input["place_type"] == "state_num_id":
        return user_input
    # Nível de região de saúde
    # health region level
    user_input["health_region_id"] = data["health_region_id"].values[0]
    if user_input["place_type"] == "health_region_id":
        return user_input
    # Nivel de cidade
    # city level
    user_input["city_id"] = data["city_id"].values[0]
    return user_input

# Escolhe Rt para SimulaCovid
# Choose Rt for SimulaCovid
def choose_rt(user_input, dfs, level):

    places = {1: "city", 2: "health_region", 3: "state"}

    place_name = places[level] + "_name"
    df = dfs[places[level]][dfs[places[level]][place_name] == user_input[place_name]]

    if level > 3:
        user_input["rt_values"] = {"best": 2.5, "worst": 3}
        user_input["rt_level"] = "nan"
        return user_input

    elif len(df["rt_low_95"].values) > 0:
        user_input["rt_values"] = {
            "best": df["rt_low_95"].values[0],
            "worst": df["rt_high_95"].values[0],
        }

        user_input["rt_level"] = [
            places[level] + "_id"
            if places[level] != "state"
            else places[level] + "_num_id"
        ][0]

        return user_input

    else:
        return choose_rt(user_input, dfs, level + 1)

# Atualiza as escolhas com a selecao do usúario
# Update choices with user selection
def update_user_input_places(user_input, dfs, config):
    # State
    # Nivel Estado
    if (
        user_input["city_name"] == "Todos"
        and user_input["health_region_name"] == "Todos"
    ):
        data = dfs["state"][dfs["state"]["state_name"] == user_input["state_name"]]
        user_input["place_type"] = "state_num_id"
        # Escolhe Rt para SimulaCovid
        user_input = choose_rt(user_input, dfs, level=3)
    # Region
    # Nivel Regional
    elif user_input["city_name"] == "Todos":
        data = dfs["health_region"][
            dfs["health_region"]["health_region_name"]
            == user_input["health_region_name"]
        ]
        user_input["place_type"] = "health_region_id"
        # Escolhe Rt para SimulaCovid
        user_input = choose_rt(user_input, dfs, level=2)
    # City
    # Nivel Cidade
    else:
        data = dfs["city"][
            (dfs["city"]["state_name"] == user_input["state_name"])
            & (dfs["city"]["city_name"] == user_input["city_name"])
        ]
        user_input["place_type"] = "city_id"
        # Escolhe Rt para SimulaCovid
        user_input = choose_rt(user_input, dfs, level=1)
    # Select location for titles
    # Seleciona localidade para títulos
    user_input["locality"] = utils.choose_place(
        city=user_input["city_name"],
        region=user_input["health_region_name"],
        state=user_input["state_name"],
    )

    # Update dos ids
    user_input = update_user_input_ids(data, user_input)
    return user_input, utils.fix_dates(data)

# Gera a tabela id "big-table" com os dados do Estado
# Generates the tabla id: "big-table" with the state data
def gen_big_table(config, dfs, currentstate):
    # st.write(dfs["state"])
    state_data = dfs["state"].sort_values(by="state_name")
    proportion = str((state_data.shape[0] + 1) * 5) + "vw"
    sector_row = state_data[state_data["state_name"] == currentstate].squeeze()
    text = f"""

FAROL COVID: Como estão os estados?

Uma visão detalhada dos indicadores para os estados brasileiros.

Estado e nível de alerta

Média móvel (últimos 7 dias) de novos casos por 100mil hab.

Taxa de contágio

Total de leitos por 100mil hab.

Taxa de subnotificação

Média móvel (últimos 7 dias) de novas mortes por 100mil hab.

Porcentagem da população vacinada

    """
    state_data = state_data[state_data["state_name"] != currentstate]
    row_order = 0
    dadosvacina = pd.read_csv('http://datasource.coronacidades.org/br/states/vacina')
    text += gen_sector_big_row(sector_row, row_order, config, dadosvacina)
    for index, sector_data in state_data.iterrows():
        row_order += 1
        text += gen_sector_big_row(sector_data, row_order, config, dadosvacina)

    text += f"""

"""
    st.write(text, unsafe_allow_html=True)

# Gera uma linha para a tabela de id "big-table" com as informacoes provenientes de uma linha de dados
# Generates a row for the table with id "big-table" with information coming from a sector data row
def gen_sector_big_row(my_state, index, config, dado):
    """ Generates a row of a table given the necessary information coming from a sector data row """
    alert_info = {
        "nan": ["grey", "❓"],
        0: ["#0090A7", "👍"],
        1: ["#F7B500", "😨"],
        2: ["#F77800", "⚠"],
        3: ["#F02C2E", "🛑"],
    }
    level_data = config["br"]["farolcovid"]["rules"]

    dado = dado[dado["state_name"] == my_state["state_name"]]
    perc_vacinados = dado.perc_vacinados.fillna("-").values[0]*2
    # TODO -> VOLTAR PROJECAO DE LEITOS
    #
{utils.dday_preffix(my_state["dday_icu_beds"])} dias

    return f"""

{my_state["state_name"]} {alert_info[my_state["overall_alert"]][1]}

{"%0.2f"%my_state["daily_cases_mavg_100k"]}

 {"%0.2f"%my_state["rt_most_likely"]}

{int(my_state["number_icu_beds_100k"])}

{int(my_state["subnotification_rate"]*100)}%

{"%0.2f"%my_state["new_deaths_mavg_100k"]}

{"%0.2f"%perc_vacinados} %

"""

def find_level(cuts, levels, value):
    if np.isnan(value):
        return "nan"
    index = bisect.bisect(cuts, value)
    return levels[min(index - 1, len(levels) - 1)]

# Puxa da API os dados
# Get data from the API
@st.cache(suppress_st_warning=True)
def get_data(config):

    dfs = {
        place: loader.read_data(
            "br",
            config,
            endpoint=config["br"]["api"]["endpoints"]["farolcovid"][place],
        ).pipe(utils.fix_dates)
        for place in ["city", "health_region", "state"]
    }

    cnes_sources = loader.read_data("br", config, "br/cities/cnes")
    # places_ids = loader.read_data("br", config, "br/places/ids")
    return dfs, cnes_sources

def main(session_state):
    # GOOGLE ANALYTICS SETUP
    if os.getenv("IS_DEV") == "FALSE":
        utils.setup_google_analytics()

    # Amplitude: Get user info
    user_analytics = amplitude.gen_user(utils.get_server_session())
    opening_response = user_analytics.safe_log_event(
        "opened farol", session_state, is_new_page=True
    )

    config = yaml.load(open("configs/config.yaml", "r"), Loader=yaml.FullLoader)

    utils.localCSS("style.css")
    st.write(
        """""",
        unsafe_allow_html=True,
    )

    utils.genHeroSection(
        title1="Farol",
        title2="Covid",
        subtitle="Entenda e controle a Covid-19 em sua cidade e estado.",
        logo="https://i.imgur.com/GLLjEYm.png",
        header=True,
        explain=True
    )

    if os.getenv("IS_HEROKU") == "TRUE":
        urlpath = os.getenv("urlpath")
    elif os.getenv("IS_DEV") == "TRUE":
        urlpath = 'http://localhost:8501/'
    else:
        urlpath = 'https://farolcovid.coronacidades.org/'
        # urlpath = 'http://localhost:8501/'
    vaccine_logo = utils.load_image("imgs/vaccine.png")
    st.write(f"""

QUER SABER MAIS SOBRE A VACINAÇÃO?

                Acompanhe nossos novos dados e descobra como avança a vacinação no seu município ou estado!

                      Veja aqui!

Saiba quando podemos controlar a pandemia no Brasil no nosso estudo realizado com dados inéditos obtidos pela Lei de Acesso à Informação.

                      Ler aqui!

""",
        unsafe_allow_html=True,
    )

    st.write(
        """

        Selecione seu estado ou município no mapa abaixo:

""",
        unsafe_allow_html=True,
    )

    # PEGA BASE DA API
    # GET DATA
    dfs, cnes_sources = get_data(config)
    map_url = "https://helper.coronacidades.org/"
    col1, col2 = st.beta_columns([0.6, 0.35])

    with col2:
        # REGIAO/CIDADE SELECAO USUARIO
        # REGION/CITY USER INPUT
        user_input = dict()
        user_input["state_name"] = st.selectbox("Estado", utils.filter_place(dfs, "state"))

        user_input["health_region_name"] = st.selectbox(
            "Região de Saúde",
            utils.filter_place(dfs, "health_region", state_name=user_input["state_name"]),
        )

        user_input["city_name"] = st.selectbox(
            "Município",
            utils.filter_place(
                dfs,
                "city",
                state_name=user_input["state_name"],
                health_region_name=user_input["health_region_name"],
            ),
        )

        changed_city = user_analytics.safe_log_event(
            "picked farol place",
            session_state,
            event_args={"state": user_input["state_name"], "city": user_input["city_name"]},
        )

        user_input, data = update_user_input_places(user_input, dfs, config)
        map_place_id = utils.Dictionary().get_state_alphabetical_id_by_name(
            user_input["state_name"]
        )
        st.write(
            f"""

        """,
            unsafe_allow_html=True,
        )
    with col1:
        # GENERATE MAPS
        # st.write(
        #     f"""
        #

        #

        #

        #

        #
        #
        #

        #

        # """,
        #     unsafe_allow_html=True,
        # )

        st.write(
            f"""

        """,
            unsafe_allow_html=True,
        )
    # st.write(
    #     f"""

    # """,
    #     unsafe_allow_html=True,
    # )

    # st.write(
    #     f"""
    #
    #
    # """,
    #     unsafe_allow_html=True,
    # )
    # st.write(
    #     f"""
    #

    #

    #
    #     """,
    #     unsafe_allow_html=True,
    # )

    # SOURCES PARAMS
    user_input = utils.get_sources(user_input, data, cnes_sources, ["beds", "icu_beds"])

    # POPULATION PARAMS
    try:
        user_input["population_params"] = {
            "N": int(data["population"].fillna(0).values[0]),
            "D": int(data["deaths"].fillna(0).values[0]),
            "I": int(data["active_cases"].fillna(0).values[0]),
            "I_confirmed": int(data["confirmed_cases"].fillna(0).values[0]),
            "I_compare": int(data["confirmed_cases"].fillna(0).values[0]),
        }
    except:
        user_input["population_params"] = {
            "N": int(data["population"].fillna(0).values[0]),
            "D": int(data["deaths"].fillna(0).values[0]),
            "I": 0,
            "I_confirmed": int(data["confirmed_cases"].fillna(0).values[0]),
            "I_compare": int(data["confirmed_cases"].fillna(0).values[0]),
        }

    user_input["Rt"] = {
        "best": data["rt_low_95"].values[0],
        "worst": data["rt_high_95"].values[0],
        "is_valid": data["rt_most_likely"].apply(str).values[0],
    }

    user_input["last_updated_cases"] = data["last_updated_subnotification"].max()
    # Update session values to standard ones if changed city or opened page or reseted values
    if (
        session_state.state_name != user_input["state_name"]
        or session_state.health_region_name != user_input["health_region_name"]
        or session_state.city_name != user_input["city_name"]
        or session_state.number_beds is None
        or session_state.reset
    ):
        session_state.state_name = user_input["state_name"]
        session_state.health_region_name = user_input["health_region_name"]
        session_state.city_name = user_input["city_name"]

        session_state.state_num_id = user_input["state_num_id"]
        session_state.health_region_id = user_input["health_region_id"]
        session_state.city_id = user_input["city_id"]

        session_state.number_beds = int(
            user_input["number_beds"]
            * config["br"]["simulacovid"]["resources_available_proportion"]
        )
        session_state.number_icu_beds = int(
            user_input["number_icu_beds"]
            * config["br"]["simulacovid"]["resources_available_proportion"]
        )
        session_state.number_cases = user_input["population_params"]["I_confirmed"]
        session_state.number_deaths = user_input["population_params"]["D"]
        session_state.reset = True

    if data["confirmed_cases"].sum() == 0:
        st.write(
            f"""

                    Seu município ou Região de Saúde ainda não possui casos reportados oficialmente. Portanto, simulamos como se o primeiro caso ocorresse hoje.

Caso queria, você pode mudar esse número abaixo:

""",
            unsafe_allow_html=True,
        )

        user_input["population_params"]["I"] = 1

    else:
        infectious_period = (
            config["br"]["seir_parameters"]["severe_duration"]
            + config["br"]["seir_parameters"]["critical_duration"]
        )
        placeholder_value_pls_solve_this = 0

    # ALERT BANNER
    # Aqui ele cria uma lista só com os estados que todas as cidades estao sem o overall_alert para criar um alerta para o estado
    states_list = dfs["city"].groupby(["state_name"]).agg({"overall_alert": "count", "state_name": "max"})
    states_list = states_list.loc[states_list['overall_alert'] < 1]
    # Adiciona forcadamente MT para a lista
    states_list = states_list.append({'state_name': 'Mato Grosso', 'overall_alert': 0}, ignore_index=True)
    states_list = states_list['state_name'].to_list()
    utils.noOverallalert(user_input, data, states_list)

    # DIMENSIONS CARDS
    dimensions = DimensionCards
    utils.genAnalysisDimmensionsSection(dimensions)

    # INDICATORS CARDS
    indicators = IndicatorCards

    indicators = update_indicators(indicators, data, config, user_input, session_state)

    data["overall_alert"] = data["overall_alert"].map(
        {0: "novo normal", 1: "moderado", 2: "alto", 3: "altíssimo"}
    )

    if "state" in user_input["place_type"]:
        # Add disclaimer to cities in state alert levels
        total_alert_regions = (
            dfs["health_region"][
                dfs["health_region"]["state_num_id"] == data["state_num_id"].unique()[0]
            ]
            .assign(
                overall_alert=lambda df: df["overall_alert"].map(
                    {0: "novo normal", 1: "moderado", 2: "alto", 3: "altíssimo"}
                )
            )["overall_alert"]
            .value_counts()
        )

        utils.genKPISection(
            place_type=user_input["place_type"],
            locality=user_input["locality"],
            alert=data["overall_alert"].values[0],
            indicators=indicators,
            n_colapse_regions=total_alert_regions[
                total_alert_regions.index.isin(["altíssimo", "alto"])
            ].sum(),
        )

    elif "city" in user_input["place_type"]:
        utils.genKPISection(
            place_type=user_input["place_type"],
            locality=user_input["locality"],
            alert=data["overall_alert"].values[0],
            indicators=indicators,
            rt_type=data["rt_place_type"].values[0],
        )

    else:
        utils.genKPISection(
            place_type=user_input["place_type"],
            locality=user_input["locality"],
            alert=data["overall_alert"].values[0],
            indicators=indicators,
        )

    # AVAILABLE CAPACITY DISCLAIMER
    # TODO -> VOLTAR PARA PROJECAO DE LEITOS
    # """
    #

    #     * Utilizamos 100% do total de leitos UTI reportados por %s em %s
    #     para cálculo da projeção de dias para atingir capacidade máxima.
Para municípios, utilizamos os recursos da respectiva regional de saúde.
    #     Leitos enfermaria contém os tipos: cirúrgicos, clínicos e hospital-dia; sendo considerado %s% já ocupado.
    #

    # """

    st.write(
        """
        <div class='base-wrapper'
            * Vacinação

“A porcentagem da população vacinada em seu local” - Total de pessoas que tomaram ao menos uma dose, dividido pelo total da população do local.

“Porcentagem da população imunizada” - Total de pessoas que receberam todas as doses recomendadas do imunizante, dividido pelo total da população do local.

“Total da população sem vacinar” - Número absoluto de habitantes do local que ainda não recebeu nenhuma dose do imunizante.

            Para mais detalhes e explicação completa confira nossa página de Metodologia no menu lateral.

        """,
        unsafe_allow_html=True,
    )

    st.write(
        """

            * Mudamos o indicador afim de refinarmos ajustes no cálculo de projeção de leitos. Entendemos que a projeção apresentada não capturava a situação da 2ª onda observada nos municípios, regionais e estados, logo substituímos este indicador por ora para revisão dos cálculos.
            Os valores de referência se baseiam nas estatísticas de países da OCDE, veja mais aqui.
            As simulações personalizadas ainda podem ser realizadas através do SimulaCovid mais abaixo.

 Leitos Enfermaria: Consideramos %s% do total reportado por %s em %s dos tipos Cirúrgico, Clínico e Hospital-dia. Para municípios, utilizamos os recursos da respectiva regional de saúde.

 Leitos UTI: Consideramos 100% do total de leitos UTI reportado por %s em %s. Para municípios, utilizamos os recursos da respectiva regional de saúde.

        """
        % (
            str(
                int(config["br"]["simulacovid"]["resources_available_proportion"] * 100)
            ),
            user_input["author_number_beds"],
            user_input["last_updated_number_beds"],
            user_input["author_number_icu_beds"], # remover na volta de projecao
            user_input["last_updated_number_icu_beds"], # remover na volta de projecao
        ),
        unsafe_allow_html=True,
    )

    # TODO: remove comment on this later!
    # utils.gen_pdf_report()

    # INDICATORS PLOTS
    if st.button("Confira a evolução de indicadores-chave"):
        opening_response = user_analytics.log_event("picked key_indicators", dict())
        if st.button("Esconder"):
            pass
        st.write(
            f"""

                    CÁLCULO DA TAXA DE CONTÁGIO EM {user_input["locality"]}

                    A taxa de contágio, conhecida como número de reprodução efetivo (Rt), traduz a dinâmica de disseminação da Covid-19 a cada dia.

O valor pode ser lido como o número médio de novas infecções diárias causadas por uma única pessoa infectada.
                    Para mais informações, visite a página de Modelos no menu lateral.

            """,
            unsafe_allow_html=True,
        )

        try:
            fig2 = plots.plot_rt_wrapper(
                user_input[user_input["place_type"]], user_input["place_type"], config
            )
            st.plotly_chart(fig2, use_container_width=True)
        except:
            st.write(
                """
Seu município, regional ou estado não possui mais de 30 dias de dados de casos confirmados.""",
                unsafe_allow_html=True,
            )
        st.write(
            "
Em breve: gráficos de subnotificação e média móvel (últimos 7 dias) de novos casos por 100k habitantes.
",
            unsafe_allow_html=True,
        )

    utils.stylizeButton(
        name="Confira a evolução de indicadores-chave",
        style_string="""border: 1px solid var(--main-white);box-sizing: border-box;border-radius: 15px; width: auto;padding: 0.5em;text-transform: uppercase;font-family: var(--main-header-font-family);color: var(--main-white);background-color: var(--main-primary);font-weight: bold;text-align: center;text-decoration: none;font-size: 18px;animation-name: fadein;animation-duration: 3s;margin-top: 1em;""",
        session_state=session_state,
    )

    # AMBASSADOR SECTION
    utils.gen_ambassador_section()

    # PDF-REPORT GEN BUTTON
    # if st.button("Gerar Relatório PDF"):
    #     user_analytics.log_event("generated pdf")
    #     st.write(
    #         """
Aguarde um momento por favor...
""",
    #         unsafe_allow_html=True,
    #     )
    #     st.markdown(
    #         pdfgen.gen_pdf_report(user_input, indicators, data, config),
    #         unsafe_allow_html=True,
    #     )

    # TOOLS
    products = ProductCards

    utils.genProductsSection(products)

    # SELECTION BUTTONS
    # TODO: limpar esse código! está 100% repetido!!!
    if session_state.continuation_selection is None:
        session_state.continuation_selection = [False, False, False, False]

    saude_button_name = "Abrir Saúde em Ordem"  # Distanciamento social
    onda_button_name = "Abrir Onda Covid "  # onda covid
    if st.button(saude_button_name):  # SAUDE
        session_state.continuation_selection = [False, False, True, False]
    if st.button(onda_button_name):  # ONDA
        session_state.continuation_selection = [False, False, False, True]

    utils.stylizeButton(
        name=saude_button_name,
        style_string="""border: 1px solid black;""",
        session_state=session_state,
        others={"ui_binSelect": 3},
    )
    utils.stylizeButton(
        name=onda_button_name,
        style_string="""border: 1px solid black;""",
        session_state=session_state,
        others={"ui_binSelect": 4},
    )
    if session_state.continuation_selection[2]:
        user_analytics.safe_log_event(
            "picked saude_em_ordem",
            session_state,
            event_args={
                "state": session_state.state_name,
                "health_region": session_state.health_region_name,
                "city": session_state.city_name,
            },
            alternatives=[
                "picked saude_em_ordem",
                "picked simulacovid",
                "picked onda",
                "picked distanciamento",
            ],
        )
        so.main(user_input, indicators, data, config, session_state)

    elif session_state.continuation_selection[3]:
        user_analytics.safe_log_event(
            "picked onda",
            session_state,
            event_args={
                "state": session_state.state_name,
                "health_region": session_state.health_region_name,
                "city": session_state.city_name,
            },
            alternatives=[
                "picked saude_em_ordem",
                "picked simulacovid",
                "picked onda",
                "picked distanciamento",
            ],
        )
        oc.main(user_input, indicators, data, config, session_state)
    st.write(
        """
        <div class='base-wrapper'
            * Desativamos Simula Covid e Distanciamento Social devido às mudanças na evolução da pandêmia.

Caso precise de alguma informação realcionada, pode nos contatar neste e-mail: contato@impulsogov.org

        """,
        unsafe_allow_html=True,
    )

    # CHAMA FUNCAO QUE GERA TABELA ID big_table
    # CALL FUNCTION TO GENERATE ID big_table
    gen_big_table(config, dfs, user_input["state_name"])
    # CHAMA FUNCOES DO UTILS PARA O FOOTER
    # CALL FUNCTIONS IN UTILS TO FOOTER
    utils.gen_footer()
    user_analytics.conclude_user_session(session_state)

if __name__ == "__main__":
    main()


# Author: Chris Moody <chrisemoody@gmail.com>
# License: MIT

# This simple example loads the newsgroups data from sklearn
# and train an LDA-like model on it
import os.path
import pickle
import time

import chainer
from chainer import cuda
from chainer import serializers
import chainer.optimizers as O
import numpy as np

from lda2vec import utils
from lda2vec import prepare_topics, print_top_words_per_topic
from lda2vec_model import LDA2Vec

gpu_id = int(os.getenv('CUDA_GPU', 0))
cuda.get_device(gpu_id).use()
print "Using GPU " + str(gpu_id)

# You must run preprocess.py before this data becomes available
vocab = pickle.load(open('../data/vocab', 'r'))
corpus = pickle.load(open('../data/corpus', 'r'))
data = np.load(open('../data/data.npz', 'r'))
flattened = data['flattened']
story_id = data['story_id']
author_id = data['author_id']
time_id = data['time_id']
ranking = data['ranking'].astype('float32')
score = data['score'].astype('float32')

# Model Parameters
# Number of documents
n_stories = story_id.max() + 1
# Number of users
n_authors = author_id.max() + 1
# Number of unique words in the vocabulary
n_vocab = flattened.max() + 1
# Number of dimensions in a single word vector
n_units = 256
# Number of topics to fit
n_story_topics = 40
n_author_topics = 20
batchsize = 4096
# Get the string representation for every compact key
words = corpus.word_list(vocab)[:n_vocab]

# How many tokens are in each story
sty_idx, lengths = np.unique(story_id, return_counts=True)
sty_len = np.zeros(sty_idx.max() + 1, dtype='int32')
sty_len[sty_idx] = lengths

# How many tokens are in each author
aut_idx, lengths = np.unique(author_id, return_counts=True)
aut_len = np.zeros(aut_idx.max() + 1, dtype='int32')
aut_len[aut_idx] = lengths

# Count all token frequencies
tok_idx, freq = np.unique(flattened, return_counts=True)
term_frequency = np.zeros(n_vocab, dtype='int32')
term_frequency[tok_idx] = freq

model = LDA2Vec(n_stories=n_stories, n_story_topics=n_story_topics,
                n_authors=n_authors, n_author_topics=n_author_topics,
                n_units=n_units, n_vocab=n_vocab, counts=term_frequency,
                n_samples=15)
if os.path.exists('lda2vec.hdf5'):
    print "Reloading from saved"
    serializers.load_hdf5("lda2vec.hdf5", model)
model.to_gpu()
optimizer = O.Adam()
optimizer.setup(model)
clip = chainer.optimizer.GradientClipping(5.0)
optimizer.add_hook(clip)

j = 0
epoch = 0
fraction = batchsize * 1.0 / flattened.shape[0]
for epoch in range(5000):
    ts = prepare_topics(cuda.to_cpu(model.mixture_sty.weights.W.data).copy(),
                        cuda.to_cpu(model.mixture_sty.factors.W.data).copy(),
                        cuda.to_cpu(model.sampler.W.data).copy(),
                        words)
    print_top_words_per_topic(ts)
    ts['doc_lengths'] = sty_len
    ts['term_frequency'] = term_frequency
    np.savez('topics.story.pyldavis', **ts)
    ta = prepare_topics(cuda.to_cpu(model.mixture_aut.weights.W.data).copy(),
                        cuda.to_cpu(model.mixture_aut.factors.W.data).copy(),
                        cuda.to_cpu(model.sampler.W.data).copy(),
                        words)
    print_top_words_per_topic(ta)
    ta['doc_lengths'] = aut_len
    ta['term_frequency'] = term_frequency
    np.savez('topics.author.pyldavis', **ta)
    for s, a, f in utils.chunks(batchsize, story_id, author_id, flattened):
        t0 = time.time()
        optimizer.zero_grads()
        l = model.fit_partial(s.copy(), a.copy(), f.copy())
        prior = model.prior()
        loss = prior * fraction
        loss.backward()
        optimizer.update()
        msg = ("J:{j:05d} E:{epoch:05d} L:{loss:1.3e} "
               "P:{prior:1.3e} R:{rate:1.3e}")
        prior.to_cpu()
        loss.to_cpu()
        t1 = time.time()
        dt = t1 - t0
        rate = batchsize / dt
        logs = dict(loss=float(l), epoch=epoch, j=j,
                    prior=float(prior.data), rate=rate)
        print msg.format(**logs)
        j += 1
    serializers.save_hdf5("lda2vec.hdf5", model)


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

from paddlespeech.t2s.audio.codec import encode_mu_law
from paddlespeech.t2s.audio.codec import float_2_label
from paddlespeech.t2s.audio.codec import label_2_float

class Clip(object):
    """Collate functor for training vocoders.
    """

    def __init__(
            self,
            batch_max_steps=20480,
            hop_size=256,
            aux_context_window=0, ):
        """Initialize customized collater for DataLoader.
        Args:

            batch_max_steps (int): The maximum length of input signal in batch.
            hop_size (int): Hop size of auxiliary features.
            aux_context_window (int): Context window size for auxiliary feature conv.

        """
        if batch_max_steps % hop_size != 0:
            batch_max_steps += -(batch_max_steps % hop_size)
        assert batch_max_steps % hop_size == 0
        self.batch_max_steps = batch_max_steps
        self.batch_max_frames = batch_max_steps // hop_size
        self.hop_size = hop_size
        self.aux_context_window = aux_context_window

        # set useful values in random cutting
        self.start_offset = aux_context_window
        self.end_offset = -(self.batch_max_frames + aux_context_window)
        self.mel_threshold = self.batch_max_frames + 2 * aux_context_window

    def __call__(self, batch):
        """Convert into batch tensors.

        Args:
            batch (list): list of tuple of the pair of audio and features. Audio shape (T, ), features shape(T', C).

        Returns:
            Tensor:
                Auxiliary feature batch (B, C, T'), where
                T = (T' - 2 * aux_context_window) * hop_size.
            Tensor:
                Target signal batch (B, 1, T).

        """
        # check length
        batch = [
            self._adjust_length(b['wave'], b['feats']) for b in batch
            if b['feats'].shape[0] > self.mel_threshold
        ]
        xs, cs = [b[0] for b in batch], [b[1] for b in batch]

        # make batch with random cut
        c_lengths = [c.shape[0] for c in cs]
        start_frames = np.array([
            np.random.randint(self.start_offset, cl + self.end_offset)
            for cl in c_lengths
        ])
        x_starts = start_frames * self.hop_size
        x_ends = x_starts + self.batch_max_steps

        c_starts = start_frames - self.aux_context_window
        c_ends = start_frames + self.batch_max_frames + self.aux_context_window
        y_batch = np.stack(
            [x[start:end] for x, start, end in zip(xs, x_starts, x_ends)])
        c_batch = np.stack(
            [c[start:end] for c, start, end in zip(cs, c_starts, c_ends)])

        # convert each batch to tensor, assume that each item in batch has the same length
        y_batch = paddle.to_tensor(
            y_batch, dtype=paddle.float32).unsqueeze(1)  # (B, 1, T)
        c_batch = paddle.to_tensor(
            c_batch, dtype=paddle.float32).transpose([0, 2, 1])  # (B, C, T')

        return y_batch, c_batch

    def _adjust_length(self, x, c):
        """Adjust the audio and feature lengths.

        Note:
            Basically we assume that the length of x and c are adjusted
            through preprocessing stage, but if we use other library processed
            features, this process will be needed.

        """
        if len(x) < c.shape[0] * self.hop_size:
            x = np.pad(x, (0, c.shape[0] * self.hop_size - len(x)), mode="edge")
        elif len(x) > c.shape[0] * self.hop_size:
            # print(
            #     f"wave length: ({len(x)}), mel length: ({c.shape[0]}), hop size: ({self.hop_size })"
            # )
            x = x[:c.shape[0] * self.hop_size]

        # check the legnth is valid
        assert len(x) == c.shape[
            0] * self.hop_size, f"wave length: ({len(x)}), mel length: ({c.shape[0]})"

        return x, c

class WaveRNNClip(Clip):
    def __init__(self,
                 mode: str='RAW',
                 batch_max_steps: int=4500,
                 hop_size: int=300,
                 aux_context_window: int=2,
                 bits: int=9,
                 mu_law: bool=True):
        self.mode = mode
        self.mel_win = batch_max_steps // hop_size + 2 * aux_context_window
        self.batch_max_steps = batch_max_steps
        self.hop_size = hop_size
        self.aux_context_window = aux_context_window
        self.mu_law = mu_law
        self.batch_max_frames = batch_max_steps // hop_size
        self.mel_threshold = self.batch_max_frames + 2 * aux_context_window
        if self.mode == 'MOL':
            self.bits = 16
        else:
            self.bits = bits

    def to_quant(self, wav):
        if self.mode == 'RAW':
            if self.mu_law:
                quant = encode_mu_law(wav, mu=2**self.bits)
            else:
                quant = float_2_label(wav, bits=self.bits)
        elif self.mode == 'MOL':
            quant = float_2_label(wav, bits=16)
        quant = quant.astype(np.int64)
        return quant

    def __call__(self, batch):
        # voc_pad = 2  this will pad the input so that the resnet can 'see' wider than input length
        # max_offsets = n_frames - 2 - (mel_win + 2 * hp.voc_pad) = n_frames - 15
        """Convert into batch tensors.
        Args:
            batch (list): list of tuple of the pair of audio and features. Audio shape (T, ), features shape(T', C).

        Returns:
            Tensor: Input signal batch (B, 1, T).
            Tensor: Target signal batch (B, 1, T).
            Tensor: Auxiliary feature batch (B, C, T'),
                where T = (T' - 2 * aux_context_window) * hop_size.

        """
        # check length
        batch = [
            self._adjust_length(b['wave'], b['feats']) for b in batch
            if b['feats'].shape[0] > self.mel_threshold
        ]
        wav, mel = [b[0] for b in batch], [b[1] for b in batch]
        # mel 此处需要转置
        mel = [x.T for x in mel]
        max_offsets = [
            x.shape[-1] - 2 - (self.mel_win + 2 * self.aux_context_window)
            for x in mel
        ]
        # the slice point of mel selecting randomly
        mel_offsets = [np.random.randint(0, offset) for offset in max_offsets]
        # the slice point of wav selecting randomly, which is behind 2(=pad) frames
        sig_offsets = [(offset + self.aux_context_window) * self.hop_size
                       for offset in mel_offsets]
        # mels.shape[1] = voc_seq_len // hop_length + 2 * voc_pad
        mels = [
            x[:, mel_offsets[i]:mel_offsets[i] + self.mel_win]
            for i, x in enumerate(mel)
        ]
        # label.shape[1] = voc_seq_len + 1
        wav = [self.to_quant(x) for x in wav]

        labels = [
            x[sig_offsets[i]:sig_offsets[i] + self.batch_max_steps + 1]
            for i, x in enumerate(wav)
        ]

        mels = np.stack(mels).astype(np.float32)
        labels = np.stack(labels).astype(np.int64)

        mels = paddle.to_tensor(mels)
        labels = paddle.to_tensor(labels, dtype='int64')
        # x is input, y is label
        x = labels[:, :self.batch_max_steps]
        y = labels[:, 1:]
        '''
        mode = RAW:
            mu_law = True:
                quant: bits = 9   0, 1, 2, ..., 509, 510, 511  int
            mu_law = False
                quant bits = 9    [0， 511]  float
        mode = MOL:
            quant: bits = 16  [0. 65536]  float
        '''
        # x should be normalizes in.[0, 1] in RAW mode
        x = label_2_float(paddle.cast(x, dtype='float32'), self.bits)
        # y should be normalizes in.[0, 1] in MOL mode
        if self.mode == 'MOL':
            y = label_2_float(paddle.cast(y, dtype='float32'), self.bits)

        return x, y, mels


import numpy as np
import matplotlib
# matplotlib.use('Agg')
import matplotlib.pyplot as plt
def get_voc_palette(num_classes):
    n = num_classes
    palette = [0]*(n*3)
    for j in range(0,n):
            lab = j
            palette[j*3+0] = 0
            palette[j*3+1] = 0
            palette[j*3+2] = 0
            i = 0
            while (lab > 0):
                    palette[j*3+0] |= (((lab >> 0) & 1) << (7-i))
                    palette[j*3+1] |= (((lab >> 1) & 1) << (7-i))
                    palette[j*3+2] |= (((lab >> 2) & 1) << (7-i))
                    i = i + 1
                    lab >>= 3
    return palette

ADE20K_palette = [0,0,0,120,120,120,180,120,120,6,230,230,80,50,50,4,200,
                    3,120,120,80,140,140,140,204,5,255,230,230,230,4,250,7,224,
                    5,255,235,255,7,150,5,61,120,120,70,8,255,51,255,6,82,143,
                    255,140,204,255,4,255,51,7,204,70,3,0,102,200,61,230,250,255,
                    6,51,11,102,255,255,7,71,255,9,224,9,7,230,220,220,220,255,9,
                    92,112,9,255,8,255,214,7,255,224,255,184,6,10,255,71,255,41,
                    10,7,255,255,224,255,8,102,8,255,255,61,6,255,194,7,255,122,8,
                    0,255,20,255,8,41,255,5,153,6,51,255,235,12,255,160,150,20,0,
                    163,255,140,140,140,250,10,15,20,255,0,31,255,0,255,31,0,255,224
                    ,0,153,255,0,0,0,255,255,71,0,0,235,255,0,173,255,31,0,255,11,200,
                    200,255,82,0,0,255,245,0,61,255,0,255,112,0,255,133,255,0,0,255,
                    163,0,255,102,0,194,255,0,0,143,255,51,255,0,0,82,255,0,255,41,0,
                    255,173,10,0,255,173,255,0,0,255,153,255,92,0,255,0,255,255,0,245,
                    255,0,102,255,173,0,255,0,20,255,184,184,0,31,255,0,255,61,0,71,255,
                    255,0,204,0,255,194,0,255,82,0,10,255,0,112,255,51,0,255,0,194,255,0,
                    122,255,0,255,163,255,153,0,0,255,10,255,112,0,143,255,0,82,0,255,163,
                    255,0,255,235,0,8,184,170,133,0,255,0,255,92,184,0,255,255,0,31,0,184,
                    255,0,214,255,255,0,112,92,255,0,0,224,255,112,224,255,70,184,160,163,
                    0,255,153,0,255,71,255,0,255,0,163,255,204,0,255,0,143,0,255,235,133,255,
                    0,255,0,235,245,0,255,255,0,122,255,245,0,10,190,212,214,255,0,0,204,255,
                    20,0,255,255,255,0,0,153,255,0,41,255,0,255,204,41,0,255,41,255,0,173,0,
                    255,0,245,255,71,0,255,122,0,255,0,255,184,0,92,255,184,255,0,0,133,255,
                    255,214,0,25,194,194,102,255,0,92,0,255]

CityScpates_palette = [128,64,128,244,35,232,70,70,70,102,102,156,190,153,153,153,153,153,
                        250,170,30,220,220,0,107,142,35,152,251,152,70,130,180,220,20,60,255,0,0,0,0,142,
                        0,0,70,0,60,100,0,80,100,0,0,230,119,11,32,128,192,0,0,64,128,128,64,128,0,192,
                        128,128,192,128,64,64,0,192,64,0,64,192,0,192,192,0,64,64,128,192,64,128,64,192,
                        128,192,192,128,0,0,64,128,0,64,0,128,64,128,128,64,0,0,192,128,0,192,0,128,192,
                        128,128,192,64,0,64,192,0,64,64,128,64,192,128,64,64,0,192,192,0,192,64,128,192,
                        192,128,192,0,64,64,128,64,64,0,192,64,128,192,64,0,64,192,128,64,192,0,192,192,
                        128,192,192,64,64,64,192,64,64,64,192,64,192,192,64,64,64,192,192,64,192,64,192,
                        192,192,192,192,32,0,0,160,0,0,32,128,0,160,128,0,32,0,128,160,0,128,32,128,128,
                        160,128,128,96,0,0,224,0,0,96,128,0,224,128,0,96,0,128,224,0,128,96,128,128,224,
                        128,128,32,64,0,160,64,0,32,192,0,160,192,0,32,64,128,160,64,128,32,192,128,160,
                        192,128,96,64,0,224,64,0,96,192,0,224,192,0,96,64,128,224,64,128,96,192,128,224,
                        192,128,32,0,64,160,0,64,32,128,64,160,128,64,32,0,192,160,0,192,32,128,192,160,
                        128,192,96,0,64,224,0,64,96,128,64,224,128,64,96,0,192,224,0,192,96,128,192,224,
                        128,192,32,64,64,160,64,64,32,192,64,160,192,64,32,64,192,160,64,192,32,192,192,
                        160,192,192,96,64,64,224,64,64,96,192,64,224,192,64,96,64,192,224,64,192,96,192,
                        192,224,192,192,0,32,0,128,32,0,0,160,0,128,160,0,0,32,128,128,32,128,0,160,128,
                        128,160,128,64,32,0,192,32,0,64,160,0,192,160,0,64,32,128,192,32,128,64,160,128,
                        192,160,128,0,96,0,128,96,0,0,224,0,128,224,0,0,96,128,128,96,128,0,224,128,128,
                        224,128,64,96,0,192,96,0,64,224,0,192,224,0,64,96,128,192,96,128,64,224,128,192,
                        224,128,0,32,64,128,32,64,0,160,64,128,160,64,0,32,192,128,32,192,0,160,192,128,
                        160,192,64,32,64,192,32,64,64,160,64,192,160,64,64,32,192,192,32,192,64,160,192,
                        192,160,192,0,96,64,128,96,64,0,224,64,128,224,64,0,96,192,128,96,192,0,224,192,
                        128,224,192,64,96,64,192,96,64,64,224,64,192,224,64,64,96,192,192,96,192,64,224,
                        192,192,224,192,32,32,0,160,32,0,32,160,0,160,160,0,32,32,128,160,32,128,32,160,
                        128,160,160,128,96,32,0,224,32,0,96,160,0,224,160,0,96,32,128,224,32,128,96,160,
                        128,224,160,128,32,96,0,160,96,0,32,224,0,160,224,0,32,96,128,160,96,128,32,224,
                        128,160,224,128,96,96,0,224,96,0,96,224,0,224,224,0,96,96,128,224,96,128,96,224,
                        128,224,224,128,32,32,64,160,32,64,32,160,64,160,160,64,32,32,192,160,32,192,32,
                        160,192,160,160,192,96,32,64,224,32,64,96,160,64,224,160,64,96,32,192,224,32,192,
                        96,160,192,224,160,192,32,96,64,160,96,64,32,224,64,160,224,64,32,96,192,160,96,
                        192,32,224,192,160,224,192,96,96,64,224,96,64,96,224,64,224,224,64,96,96,192,224,
                        96,192,96,224,192,0,0,0]

COCO_palette = [31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227,
    119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44,
    214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207,
    31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75,
    227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44,
    214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189,
    34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75,
    227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127,
    14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189,
    34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103,
    189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127,
    14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127
    , 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103,
    189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14,
    44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127,
    127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189,
    140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44,
    160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190,
    207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194,
    127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148,
    103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127,
    14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34,
    23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227,
    119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14, 44, 160, 44, 214, 39,
    40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127, 127, 188, 189, 34, 23, 190, 207, 31, 119,
    180, 255, 127, 14, 44, 160, 44, 214, 39, 40, 148, 103, 189, 140, 86, 75, 227, 119, 194, 127, 127,
    127, 188, 189, 34, 23, 190, 207, 31, 119, 180, 255, 127, 14]

def heatmap(data, row_labels, col_labels, ax=None,
            cbar_kw={}, cbarlabel="", **kwargs):
    """
    Create a heatmap from a numpy array and two lists of labels.

    Parameters
    ----------
    data
        A 2D numpy array of shape (N, M).
    row_labels
        A list or array of length N with the labels for the rows.
    col_labels
        A list or array of length M with the labels for the columns.
    ax
        A `matplotlib.axes.Axes` instance to which the heatmap is plotted.  If
        not provided, use current axes or create a new one.  Optional.
    cbar_kw
        A dictionary with arguments to `matplotlib.Figure.colorbar`.  Optional.
    cbarlabel
        The label for the colorbar.  Optional.
    **kwargs
        All other arguments are forwarded to `imshow`.
    """

    if not ax:
        ax = plt.gca()

    # Plot the heatmap
    im = ax.imshow(data, **kwargs)

    # Create colorbar
    cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw)
    cbar.ax.set_ylabel(cbarlabel, rotation=-90, va="bottom")

    # We want to show all ticks...
    ax.set_xticks(np.arange(data.shape[1]))
    ax.set_yticks(np.arange(data.shape[0]))
    # ... and label them with the respective list entries.
    ax.set_xticklabels(col_labels)
    ax.set_yticklabels(row_labels)

    # Let the horizontal axes labeling appear on top.
    ax.tick_params(top=True, bottom=False,
                   labeltop=True, labelbottom=False)

    # Rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=-30, ha="right",
             rotation_mode="anchor")

    # Turn spines off and create white grid.
    for edge, spine in ax.spines.items():
        spine.set_visible(False)

    ax.set_xticks(np.arange(data.shape[1]+1)-.5, minor=True)
    ax.set_yticks(np.arange(data.shape[0]+1)-.5, minor=True)
    ax.grid(which="minor", color="w", linestyle='-', linewidth=3)
    ax.tick_params(which="minor", bottom=False, left=False)

    return im, cbar

def annotate_heatmap(im, data=None, valfmt="{x:.2f}",
                     textcolors=["black", "white"],
                     threshold=None, **textkw):
    """
    A function to annotate a heatmap.

    Parameters
    ----------
    im
        The AxesImage to be labeled.
    data
        Data used to annotate.  If None, the image's data is used.  Optional.
    valfmt
        The format of the annotations inside the heatmap.  This should either
        use the string format method, e.g. "$ {x:.2f}", or be a
        `matplotlib.ticker.Formatter`.  Optional.
    textcolors
        A list or array of two color specifications.  The first is used for
        values below a threshold, the second for those above.  Optional.
    threshold
        Value in data units according to which the colors from textcolors are
        applied.  If None (the default) uses the middle of the colormap as
        separation.  Optional.
    **kwargs
        All other arguments are forwarded to each call to `text` used to create
        the text labels.
    """

    if not isinstance(data, (list, np.ndarray)):
        data = im.get_array()

    # Normalize the threshold to the images color range.
    if threshold is not None:
        threshold = im.norm(threshold)
    else:
        threshold = im.norm(data.max())/2.

    # Set default alignment to center, but allow it to be
    # overwritten by textkw.
    kw = dict(horizontalalignment="center",
              verticalalignment="center")
    kw.update(textkw)

    # Get the formatter in case a string is supplied
    if isinstance(valfmt, str):
        valfmt = matplotlib.ticker.StrMethodFormatter(valfmt)

    # Loop over the data and create a `Text` for each "pixel".
    # Change the text's color depending on the data.
    texts = []
    for i in range(data.shape[0]):
        for j in range(data.shape[1]):
            kw.update(color=textcolors[int(im.norm(data[i, j]) > threshold)])
            text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)
            texts.append(text)

    return texts

def save_confusion_matrix(array):
    fig, ax = plt.subplots()
    name =['road','sidewalk','building','wall','fence','pole','traffic light','traffic sign',
           'vegetation','terrain','sky', 'person','rider','car','truck','bus','train',
           'motorcycle','bicycle']
    im, cbar = heatmap(array,name, name, ax=ax,
                       cmap="YlGn", cbarlabel="harvest [t/year]")
    texts = annotate_heatmap(im, valfmt="{x:.3f} t")

    fig.tight_layout()
    plt.savefig('cm.png', dpi=100)


# -*- coding: utf-8 -*-
import os
import pandas as pd
from multiprocessing import Pool
import warnings
import numpy as np
import scipy
import pybacktest as pb
import matplotlib.pyplot as plt
import threading
import multiprocessing
import tensorflow as tf
import tensorflow.contrib.slim as slim
import scipy.signal
from random import choice
from time import sleep
from time import time
import sys
from trader_gym import environment
from A3C_class import *
from configs import TRAIN_DATA, LOAD_MODEL, LR, FRAMES_STACKED, NUM_WORKERS, MODEL_DIR
warnings.filterwarnings("ignore")

train_df = pd.read_pickle(TRAIN_DATA)

if FRAMES_STACKED > 1:
    data = np.hstack([train_df.values[i:-FRAMES_STACKED + i - 1, :] for i in range(FRAMES_STACKED, 0, -1)])
else:
    data = train_df.values

train_df = pd.DataFrame(data, train_df[FRAMES_STACKED:-1].index)
max_episode_len = train_df.shape[0]

tf.reset_default_graph()
if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR)

with tf.device(USE_DEVICE):
    global_episodes = tf.Variable(0, dtype=tf.int32, name='global_episodes', trainable=False)
    trainer = tf.train.RMSPropOptimizer(learning_rate=LR, decay=0.99, epsilon=1e-6)
    master_network = AC_Network(s_size, a_size, 'global', None)
    workers = []
    for i in range(NUM_WORKERS):
        env = environment(train_df, max_episode_len)
        workers.append(Worker(env, i, s_size, a_size, trainer, MODEL_DIR, global_episodes))
    saver = tf.train.Saver(max_to_keep=25)

if 'session' in locals() and session is not None:
    print('Close interactive session')
    session.close()
if 'sess' in locals() and sess is not None:
    print('Close interactive session')
    sess.close()

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0)

with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True, gpu_options=gpu_options)) as sess:
    coord = tf.train.Coordinator()
    if LOAD_MODEL:
        print('Loading Model...')
        ckpt = tf.train.get_checkpoint_state(MODEL_DIR)
        saver.restore(sess, ckpt.model_checkpoint_path)
    else:
        sess.run(tf.global_variables_initializer())
    summary_writer = tf.summary.FileWriter("tb/train", graph=sess.graph)
    worker_threads = []
    for worker in workers:
        def worker_work(): return worker.work(max_episode_len, gamma, sess, coord, saver, FRAMES_STACKED)
        t = threading.Thread(target=(worker_work))
        t.start()
        worker_threads.append(t)
    coord.join(worker_threads)


import numpy as np
import os
from brightway2 import projects, methods, Method
import pyprind
from .utils import _check_project, _check_result_dir, _check_method
from pathlib import Path
import json

def concat_arrays(l):
    """Concatenate column-wise arrays in list

    Parameters
    ----------
    l : list
        List of paths to ndarray

    Returns
    -------
    Single concatenated array
    """
    arrs = [np.load(str(arr)) for arr in l]
    if not len(set([arr.shape[0] for arr in arrs])):
        raise ValueError(
            "Arrays don't all have the same number of rows for act {}".format(
                l[0].name
            )
        )
    return np.concatenate(arrs, axis=1)

def find_common_act_codes(sb_dirpaths):
    """Identify activitity codes common to samples_batches

    Parameters
    ----------
    sb_dirpaths : list
        List of paths to samples_batch directories

    Returns
    -------
    common_act_codes, all_sb_dirs_equal
    """
    available_act_codes = [
        set([f[0:-4] for f in os.listdir(fp) if f[-4:] == '.npy'])
        for fp in sb_dirpaths
    ]

    common_act_codes = set.intersection(*available_act_codes)
    all_same = all([common_act_codes==avail_sb for avail_sb in available_act_codes])
    return list(common_act_codes), all_same

def default_sim_name(sb_dirpaths):
    """Generate a default sim_name

    Parameters
    ----------
    sb_dirpaths : list
        List of paths to samples_batch directories

    Returns
    -------
    sim_name : str
    """
    return "_".join([Path(dirpath).name for dirpath in sb_dirpaths])

def concat_samples_arrays_in_result_type_dir(
        result_dir, sb_id_list=None, result_type_dirname="LCI",
        totals_or_per_exchanges=None, sim_name=None, dest=None,
        fail_if_samples_batches_different=False, ignore_missing=True):
    """ Concatenate and save LCI sample arrays in result dir

    Parameters
    ------------
    result_type : str
        Path to the directory where data used or generated by bw2preagg is saved.
    sb_id_list : list, default=None
        List of samples_batch ids. If None, all samples_batch are processed.
    result_type_dirname : str, default="LCI"
        Name of the result type subdirectory where samples batches are found
    sim_name : str, default=None
        Name to give the directory in which results will be saved
    dest : str
        Path to location where concatenated array directory will be saved
    fail_if_samples_batches_different : bool, default=False
        If False, will raise ValueError if arrays available in samples_batch
        folders are not the same
        If True, will concatenate only those arrays that are available in all
        samples_batch folders
    ignore_missing : bool, default=True
        If False, will concatenate arrays in LCI folder only if all expected
        LCI arrays are present
    """
    result_dir = Path(_check_result_dir(result_dir))
    if result_type_dirname == "LCI":
        result_type_dir = result_dir / "probabilistic" / "LCI"
    else:
        result_type_dir = result_dir / "probabilistic" / result_type_dirname / totals_or_per_exchanges
    if not result_type_dir.is_dir():
        raise ValueError("No arrays available at {}".format(result_type_dir))
    available_sb = os.listdir(result_type_dir)
    if sb_id_list is not None:
        missing_sb = [sb for sb in sb_id_list if sb not in available_sb]
        if missing_sb:
            raise ValueError("Not all specified sample batches are available in {}. Missing: {}".format(result_type_dir, missing_sb))
    else:
        sb_id_list = sorted(available_sb)
    sb_fp_list = [result_type_dir / sb_id for sb_id in sb_id_list]
    available_acts, all_dirs_equal = find_common_act_codes(sb_fp_list)
    if not all_dirs_equal and fail_if_samples_batches_different:
        raise ValueError("Samples batches did not contain arrays for the same activities.")
    if ignore_missing is False:
        with open(result_dir / "common_files" / "ordered_activity_codes.json", "r") as f:
            expected_act_codes = json.load(f)
        if set(available_acts) != set(expected_act_codes):
            raise ValueError("The expected LCI arrays and the common LCI arrays in the LCI folders differ.")
    if not sim_name:
        sim_name = default_sim_name(sb_fp_list)
    if not dest:
        dest = result_dir / "concatenated"
    save_dir = Path(dest) / result_type_dirname / sim_name
    save_dir.mkdir(exist_ok=True, parents=True)
    for act_code in available_acts:
        concat_lci_array = concat_arrays(
            [
                dirpath / "{}.npy".format(act_code)
                for dirpath in sb_fp_list
            ]
        )
        np.save(str(save_dir / "{}.npy".format(act_code)), concat_lci_array)
    return

def concat_lcia_samples_arrays_from_method_tuple(
        result_dir, method, sb_id_list=None, totals_or_per_exchanges="totals",
        project_name=None, sim_name=None, dest=None,
        fail_if_samples_batches_different=False, ignore_missing=True):
    """ Concatenate and save LCIA sample arrays in result dir

    Parameters
    ------------
    result_type : str
        Path to the directory where data used or generated by bw2preagg is saved.
    method : tuple
        LCIA method identification in brightway2 (tuple)
    sb_id_list : list
        List of samples_batch ids. If None, all samples_batch are processed.
    totals_or_per_exchanges : str, default="totals"
        Deal with totals or results per elementary flows
    project_name : str
        Name of the brightway2 project where the database is imported
    sim_name : str, default=None
        Name to give the directory in which results will be saved. If None, a
        default name is generated from sb_id_list
    dest : str
        Path to location where concatenated array directory will be saved
    fail_if_samples_batches_different : bool, default=False
        If False, will raise ValueError if arrays available in samples_batch
        folders are not the same
        If True, will concatenate only those arrays that are available in all
        samples_batch folders
    ignore_missing : bool, default=True
        If False, will concatenate arrays in LCI folder only if all expected
        LCI arrays are present
    """
    if project_name:
        projects.set_current(_check_project(project_name))
    method = _check_method(method)
    abbr = Method(method).get_abbreviation()
    concat_samples_arrays_in_result_type_dir(
        result_dir=result_dir, sb_id_list=sb_id_list,
        result_type_dirname=abbr, totals_or_per_exchanges=totals_or_per_exchanges,
        sim_name=sim_name, dest=dest,
        fail_if_samples_batches_different=fail_if_samples_batches_different,
        ignore_missing=ignore_missing)


import numpy as np
import keras
from keras.optimizers import Adam
from nets.efficientdet_training import Generator
from nets.efficientdet_training import focal,smooth_l1
from nets.efficientdet import Efficientdet
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
from utils.utils import BBoxUtility
from utils.anchors import get_anchors

#---------------------------------------------------#
#   获得类和先验框
#---------------------------------------------------#
def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

freeze_layers = [226, 328, 328, 373, 463, 565, 655, 802]
image_sizes = [512, 640, 768, 896, 1024, 1280, 1408, 1536]

#----------------------------------------------------#
#   检测精度mAP和pr曲线计算参考视频
#   https://www.bilibili.com/video/BV1zE411u7Vw
#----------------------------------------------------#
if __name__ == "__main__":
    #-------------------------------------------#
    #   训练前，请指定好phi和model_path
    #   二者所使用Efficientdet版本要相同
    #-------------------------------------------#
    phi = 0
    annotation_path = '2007_train.txt'

    classes_path = 'model_data/voc_classes.txt'
    class_names = get_classes(classes_path)
    NUM_CLASSES = len(class_names)

    #------------------------------------------------------#
    #   权值文件请看README，百度网盘下载
    #   训练自己的数据集时提示维度不匹配正常
    #   预测的东西都不一样了自然维度不匹配
    #------------------------------------------------------#
    model_path = "model_data/efficientdet-d0-voc.h5"

    model = Efficientdet(phi,num_classes=NUM_CLASSES)
    priors = get_anchors(image_sizes[phi])
    bbox_util = BBoxUtility(NUM_CLASSES, priors)

    model.load_weights(model_path,by_name=True,skip_mismatch=True)

    # 0.1用于验证，0.9用于训练
    val_split = 0.1
    with open(annotation_path) as f:
        lines = f.readlines()
    np.random.seed(10101)
    np.random.shuffle(lines)
    np.random.seed(None)
    num_val = int(len(lines)*val_split)
    num_train = len(lines) - num_val

    # 训练参数设置
    logging = TensorBoard(log_dir="logs")
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=2, verbose=1)
    checkpoint = ModelCheckpoint('logs/ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
        monitor='val_loss', save_weights_only=True, save_best_only=False, period=1)
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=6, verbose=1)

    for i in range(freeze_layers[phi]):
        model.layers[i].trainable = False

    #------------------------------------------------------#
    #   主干特征提取网络特征通用，冻结训练可以加快训练速度
    #   也可以在训练初期防止权值被破坏。
    #   Init_Epoch为起始世代
    #   Freeze_Epoch为冻结训练的世代
    #   Epoch总训练世代
    #   提示OOM或者显存不足请调小Batch_size
    #------------------------------------------------------#
    if True:
        #--------------------------------------------#
        #   BATCH_SIZE不要太小，不然训练效果很差
        #--------------------------------------------#
        BATCH_SIZE = 4
        Lr = 1e-3
        Init_Epoch = 0
        Freeze_Epoch = 50
        gen = Generator(bbox_util, BATCH_SIZE, lines[:num_train], lines[num_train:],
                        (image_sizes[phi], image_sizes[phi]),NUM_CLASSES)
        model.compile(loss={
                    'regression'    : smooth_l1(),
                    'classification': focal()
                },optimizer=keras.optimizers.Adam(Lr)
        )
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, BATCH_SIZE))
        model.fit_generator(
                gen.generate(True),
                steps_per_epoch=max(1, num_train//BATCH_SIZE),
                validation_data=gen.generate(False),
                validation_steps=max(1, num_val//BATCH_SIZE),
                epochs=Freeze_Epoch,
                verbose=1,
                initial_epoch=Init_Epoch ,
                callbacks=[logging, checkpoint, reduce_lr, early_stopping]
            )

    for i in range(freeze_layers[phi]):
        model.layers[i].trainable = True

    if True:
        #--------------------------------------------#
        #   BATCH_SIZE不要太小，不然训练效果很差
        #--------------------------------------------#
        BATCH_SIZE = 4
        Lr = 5e-5
        Freeze_Epoch = 50
        Epoch = 100
        gen = Generator(bbox_util, BATCH_SIZE, lines[:num_train], lines[num_train:],
                        (image_sizes[phi], image_sizes[phi]),NUM_CLASSES)

        model.compile(loss={
                    'regression'    : smooth_l1(),
                    'classification': focal()
                },optimizer=keras.optimizers.Adam(Lr)
        )
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, BATCH_SIZE))
        model.fit_generator(
                gen.generate(True),
                steps_per_epoch=max(1, num_train//BATCH_SIZE),
                validation_data=gen.generate(False),
                validation_steps=max(1, num_val//BATCH_SIZE),
                epochs=Epoch,
                verbose=1,
                initial_epoch=Freeze_Epoch,
                callbacks=[logging, checkpoint, reduce_lr, early_stopping]
            )


#
import unittest
import numpy as np
import tensorflow as tf
from app.drl.probability_distribution import ProbabilityDistribution

class TProbabilityDistribution(unittest.TestCase):
    def test_call(self):
        print('测试ProbabilityDistribution...')
        pd = ProbabilityDistribution()
        logits = tf.log([[1., 900., 1., 500., 1., 1., 1.]])
        for i in range(10):
            a1 = pd.predict(logits, steps=1)
            print('action={0}'.format(a1[0]))
        self.assertTrue(True)


"""
Docks protein-ligand pairs
"""
__author__ = "Bharath Ramsundar"
__copyright__ = "Copyright 2016, Stanford University"
__license__ = "MIT"

import logging
import numpy as np
import os
import tempfile
from deepchem.data import DiskDataset
from deepchem.models import SklearnModel
from deepchem.models import MultitaskRegressor
from deepchem.dock.pose_scoring import GridPoseScorer
from deepchem.dock.pose_generation import VinaPoseGenerator
from sklearn.ensemble import RandomForestRegressor
from subprocess import call

logger = logging.getLogger(__name__)

class Docker(object):
  """Abstract Class specifying API for Docking."""

  def dock(self,
           protein_file,
           ligand_file,
           centroid=None,
           box_dims=None,
           dry_run=False):
    raise NotImplementedError

class VinaGridRFDocker(Docker):
  """Vina pose-generation, RF-models on grid-featurization of complexes."""

  def __init__(self, exhaustiveness=10, detect_pockets=False):
    """Builds model."""
    self.base_dir = tempfile.mkdtemp()
    logger.info("About to download trained model.")
    call((
        "wget -nv -c http://deepchem.io.s3-website-us-west-1.amazonaws.com/trained_models/random_full_RF.tar.gz"
    ).split())
    call(("tar -zxvf random_full_RF.tar.gz").split())
    call(("mv random_full_RF %s" % (self.base_dir)).split())
    self.model_dir = os.path.join(self.base_dir, "random_full_RF")

    # Fit model on dataset
    model = SklearnModel(model_dir=self.model_dir)
    model.reload()

    self.pose_scorer = GridPoseScorer(model, feat="grid")
    self.pose_generator = VinaPoseGenerator(
        exhaustiveness=exhaustiveness, detect_pockets=detect_pockets)

  def dock(self,
           protein_file,
           ligand_file,
           centroid=None,
           box_dims=None,
           dry_run=False):
    """Docks using Vina and RF."""
    protein_docked, ligand_docked = self.pose_generator.generate_poses(
        protein_file, ligand_file, centroid, box_dims, dry_run)
    if not dry_run:
      score = self.pose_scorer.score(protein_docked, ligand_docked)
    else:
      score = np.zeros((1,))
    return (score, (protein_docked, ligand_docked))

'''
class VinaGridDNNDocker(object):
  """Vina pose-generation, DNN-models on grid-featurization of complexes."""

  def __init__(self, exhaustiveness=10, detect_pockets=False):
    """Builds model."""
    self.base_dir = tempfile.mkdtemp()
    logger.info("About to download trained model.")
    call((
        "wget -nv -c http://deepchem.io.s3-website-us-west-1.amazonaws.com/trained_models/random_full_DNN.tar.gz"
    ).split())
    call(("tar -zxvf random_full_DNN.tar.gz").split())
    call(("mv random_full_DNN %s" % (self.base_dir)).split())
    self.model_dir = os.path.join(self.base_dir, "random_full_DNN")

    # Fit model on dataset
    pdbbind_tasks = ["-logKd/Ki"]
    n_features = 2052
    model = MultitaskRegressor(
        len(pdbbind_tasks),
        n_features,
        dropouts=[.25],
        learning_rate=0.0003,
        weight_init_stddevs=[.1],
        batch_size=64,
        model_dir=self.model_dir)
    model.reload()

    self.pose_scorer = GridPoseScorer(model, feat="grid")
    self.pose_generator = VinaPoseGenerator(
        exhaustiveness=exhaustiveness, detect_pockets=detect_pockets)

  def dock(self,
           protein_file,
           ligand_file,
           centroid=None,
           box_dims=None,
           dry_run=False):
    """Docks using Vina and DNNs."""
    protein_docked, ligand_docked = self.pose_generator.generate_poses(
        protein_file, ligand_file, centroid, box_dims, dry_run)
    if not dry_run:
      score = self.pose_scorer.score(protein_docked, ligand_docked)
    else:
      score = np.zeros((1,))
    return (score, (protein_docked, ligand_docked))
'''


# import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

import time
import numpy as np

class TensorFaceDetector(object):
    def __init__(self, path_to_ckpt=None):
        print("init")
        if path_to_ckpt is None:
            path_to_ckpt = '/home/kom/code/men/mayaclient/models/frozen_inference_graph_face.pb'

        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(path_to_ckpt, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        with self.detection_graph.as_default():
            config = tf.ConfigProto()
            # config.gpu_options.per_process_gpu_memory_fraction=0.6
            # config.intra_op_parallelism_threads=2
            # config.inter_op_parallelism_threads=2
            config.gpu_options.allow_growth = True
            self.sess = tf.Session(graph=self.detection_graph, config=config)
            self.windowNotSet = True

    def __call__(self):
        print('call')

    def __del__(self):
        self.sess.close()

    def run(self, image):
        image_np = image
        # image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # the array based representation of the image will be used later in order to prepare the
        # result image with boxes and labels on it.
        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
        image_np_expanded = np.expand_dims(image_np, axis=0)
        image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
        classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
        num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
        # Actual detection.
        start_time = time.time()
        (boxes, scores, classes, num_detections) = self.sess.run(
            [boxes, scores, classes, num_detections],
            feed_dict={image_tensor: image_np_expanded})
        elapsed_time = time.time() - start_time
        print('inference time cost: {}'.format(elapsed_time))

        return (boxes, scores, classes, num_detections)


#################### Main module which connects all the others together ####################
# import necessary packages
import Eye_locator
import controller

import time

import cv2
import numpy as np
import dlib

import tensorflow.compat.v1 as tf

import logger
import config
import inferno
import utils

# Create a dlib face detector
detector = dlib.get_frontal_face_detector()
# Create a 81 points based landmark detector
predictor = dlib.shape_predictor('shape_predictor_81_face_landmarks.dat')
# Create a Haar Classifer for the eye
haar_classifier = cv2.CascadeClassifier('haarcascade_eye_tree_eyeglasses.xml')

# Set the initial eye to be chosen as left
direction = 'left'

ret = True
# Create cv2 video capture for capturing the video
cap = cv2.VideoCapture(0);
# read the first frame
_,k = cap.read()

# get the size of the frame and set the desired position as center of the frame
des_pos = [k.shape[1]//2, k.shape[0]//2]

# Set the PID constants for ech axis
# Need to be tuned based on hardware
pid_x = [0.1,0,0.1]
pid_y = [0.1,0,0.1]
# Store the error data of previous interation for next one
# prev_error_data -> [error_x, error_y, sum_error_x, sum_error_y]
prev_error_data = [0, 0, 0, 0]

# Create a logger
logger = logger.Logger('INC', "3A4Bh-Ref25","", config.config, dir="models/")
logger.log("Start inferring model...")
# Create a tensorflow session
sess = tf.Session()
# Load the model
model = inferno.load_model(sess, 'INC', '3A4Bh-Ref25', logger)

# reset the position of the x-y-z platform
controller.reset(cap, detector)

# Infinite loop
while ret:
    ret, frame = cap.read() # Read the frame from the camera

    # Get the eye location by running the function in module Eye_locator
    '''
        x -> x location of the center of the eyeball/pupil
        y -> y location of the center of the eyeball/pupil
        r -> radius of the eyeball/pupil
        eye -> segmented out image of the eye
    '''
    x, y, r, eye = Eye_locator.get_eye_location(frame, direction, detector, predictor, haar_classifier ,sess, model, logger)
    # If eye detection failed, then it is assumed that the camera is very close to the eye
    if x == -1 or y == -1 or r == -1:
        # If direction being detected is left, then change it to right, move the camera a bit forward and take an Image
        # This is the image of the eye, store it
        if direction == 'left':
            ## move forward in z by a small distance
            _,f = cap.read()
            cv2.imwrite('left.jpg',f)
            direction = 'right'
            # reset the position of the x-y-z platform
            controller.reset(cap,detector)
            # Set the error data to 0
            prev_error_data = [0, 0, 0, 0]
            continue
        # if the direction being detected is right, then move the camera a bit forward and take an Image
        # this is the image of the right eye, store it and exit the program
        elif direction == 'right':
            ## move forward in z by a small distance
            _,f = cap.read()
            cv2.imwrite('right.jpg', f)
            break
        continue

    # Get the distance to be moved from the PID controller, given the current and desired location of the eye
    dist_x, dist_y, prev_error_data = controller.calc_distance([x,y], des_pos, pid_x, pid_y, prev_error_data)

    # Calculate the number of steps to rotate the stepper motor by to move by the given distance
    num_steps_x = controller.calc_steps(dist_x)
    num_steps_y = controller.calc_steps(dist_y)

    # If the distance to be moved is less than some threshold which needs to be tuned based on the hardware, move forward
    if abs(num_steps_x) < 100 and abs(num_steps_y) < 100:
        print('Move forward')
        # send a fixed forward motion number of steps
    else:
        # move the camera in appropriate direction to attain the desired position
        if num_steps_x < 0:
            print('Move right')
            # send this value to the micro-controller
        else:
            print('Move left')
            # send this value to the micro-controller

        if num_steps_y < 0:
            print('Move down')
            # send this value to the micro-controller
        else:
            print('Move up')
            # send this value to the micro-controller

    # Show the eyeball/pupil on the frame captured
    frame = utils.annotator((0,255,0), frame, x, y, r)
    # Display the frame
    cv2.imshow("Image",frame)

    # if q is pressed, exit the program
    key = cv2.waitKey(1) & 0xFF
    if key==ord('q'):
        break
# perform cleanup
cv2.destroyAllWindows()
cap.release()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 31 12:20:39 2021

@author: f004swn
"""
import os
import numpy as np
from keras.layers import Dropout, ZeroPadding2D, Convolution2D, MaxPooling2D, Flatten, Activation
from keras import Sequential
from keras.applications.resnet import preprocess_input
import cv2
from tensorflow import keras
from mrcnn.model import MaskRCNN
from mrcnn.config import Config

os.chdir('/Users/f004swn/Documents/Code/packages')

modelfile = 'psypose/models/model_weights/deepfashion2.h5'

class_names = ['short_sleeved_shirt', 'long_sleeved_shirt', 'short_sleeved_outwear', 'long_sleeved_outwear', 'vest', 'sling',
               'shorts', 'trousers', 'skirt', 'short_sleeved_dress', 'long_sleeved_dress',
               'vest_dress', 'sling_dress']

class TestConfig(Config):
     NAME = "test"
     GPU_COUNT = 1
     IMAGES_PER_GPU = 1
     NUM_CLASSES = 1 + 13

rcnn = MaskRCNN(mode='inference', config=TestConfig())
rcnn.load_weights(modelfile, by_name=True)


import os
import datetime
import hashlib
import multiprocessing as mp
import queue
from multiprocessing.connection import Connection
from abc import ABC, abstractmethod
from typing import Dict
import numpy as np
import tflite_runtime.interpreter as tflite
from tflite_runtime.interpreter import load_delegate
from frigate.util import EventsPerSecond, listen, SharedMemoryFrameManager

def load_labels(path, encoding='utf-8'):
  """Loads labels from file (with or without index numbers).
  Args:
    path: path to label file.
    encoding: label file encoding.
  Returns:
    Dictionary mapping indices to labels.
  """
  with open(path, 'r', encoding=encoding) as f:
    lines = f.readlines()
    if not lines:
        return {}

    if lines[0].split(' ', maxsplit=1)[0].isdigit():
        pairs = [line.split(' ', maxsplit=1) for line in lines]
        return {int(index): label.strip() for index, label in pairs}
    else:
        return {index: line.strip() for index, line in enumerate(lines)}

class ObjectDetector(ABC):
    @abstractmethod
    def detect(self, tensor_input, threshold = .4):
        pass

class LocalObjectDetector(ObjectDetector):
    def __init__(self, tf_device=None, labels=None):
        self.fps = EventsPerSecond()
        if labels is None:
            self.labels = {}
        else:
            self.labels = load_labels(labels)

        device_config = {"device": "usb"}
        if not tf_device is None:
            device_config = {"device": tf_device}

        edge_tpu_delegate = None

        if tf_device != 'cpu':
            try:
                print(f"Attempting to load TPU as {device_config['device']}")
                edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', device_config)
                print("TPU found")
            except ValueError:
                print("No EdgeTPU detected. Falling back to CPU.")

        if edge_tpu_delegate is None:
            self.interpreter = tflite.Interpreter(
                model_path='/cpu_model.tflite')
        else:
            self.interpreter = tflite.Interpreter(
                model_path='/edgetpu_model.tflite',
                experimental_delegates=[edge_tpu_delegate])

        self.interpreter.allocate_tensors()

        self.tensor_input_details = self.interpreter.get_input_details()
        self.tensor_output_details = self.interpreter.get_output_details()

    def detect(self, tensor_input, threshold=.4):
        detections = []

        raw_detections = self.detect_raw(tensor_input)

        for d in raw_detections:
            if d[1] < threshold:
                break
            detections.append((
                self.labels[int(d[0])],
                float(d[1]),
                (d[2], d[3], d[4], d[5])
            ))
        self.fps.update()
        return detections

    def detect_raw(self, tensor_input):
        self.interpreter.set_tensor(self.tensor_input_details[0]['index'], tensor_input)
        self.interpreter.invoke()
        boxes = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[0]['index']))
        label_codes = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[1]['index']))
        scores = np.squeeze(self.interpreter.get_tensor(self.tensor_output_details[2]['index']))

        detections = np.zeros((20,6), np.float32)
        for i, score in enumerate(scores):
            detections[i] = [label_codes[i], score, boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]]

        return detections

def run_detector(detection_queue, out_events: Dict[str, mp.Event], avg_speed, start, tf_device):
    print(f"Starting detection process: {os.getpid()}")
    listen()
    frame_manager = SharedMemoryFrameManager()
    object_detector = LocalObjectDetector(tf_device=tf_device)

    outputs = {}
    for name in out_events.keys():
        out_shm = mp.shared_memory.SharedMemory(name=f"out-{name}", create=False)
        out_np = np.ndarray((20,6), dtype=np.float32, buffer=out_shm.buf)
        outputs[name] = {
            'shm': out_shm,
            'np': out_np
        }

    while True:
        connection_id = detection_queue.get()
        input_frame = frame_manager.get(connection_id, (1,300,300,3))

        if input_frame is None:
            continue

        # detect and send the output
        start.value = datetime.datetime.now().timestamp()
        detections = object_detector.detect_raw(input_frame)
        duration = datetime.datetime.now().timestamp()-start.value
        outputs[connection_id]['np'][:] = detections[:]
        out_events[connection_id].set()
        start.value = 0.0

        avg_speed.value = (avg_speed.value*9 + duration)/10

class EdgeTPUProcess():
    def __init__(self, detection_queue, out_events, tf_device=None):
        self.out_events = out_events
        self.detection_queue = detection_queue
        self.avg_inference_speed = mp.Value('d', 0.01)
        self.detection_start = mp.Value('d', 0.0)
        self.detect_process = None
        self.tf_device = tf_device
        self.start_or_restart()

    def stop(self):
        self.detect_process.terminate()
        print("Waiting for detection process to exit gracefully...")
        self.detect_process.join(timeout=30)
        if self.detect_process.exitcode is None:
            print("Detection process didnt exit. Force killing...")
            self.detect_process.kill()
            self.detect_process.join()

    def start_or_restart(self):
        self.detection_start.value = 0.0
        if (not self.detect_process is None) and self.detect_process.is_alive():
            self.stop()
        self.detect_process = mp.Process(target=run_detector, args=(self.detection_queue, self.out_events, self.avg_inference_speed, self.detection_start, self.tf_device))
        self.detect_process.daemon = True
        self.detect_process.start()

class RemoteObjectDetector():
    def __init__(self, name, labels, detection_queue, event):
        self.labels = load_labels(labels)
        self.name = name
        self.fps = EventsPerSecond()
        self.detection_queue = detection_queue
        self.event = event
        self.shm = mp.shared_memory.SharedMemory(name=self.name, create=False)
        self.np_shm = np.ndarray((1,300,300,3), dtype=np.uint8, buffer=self.shm.buf)
        self.out_shm = mp.shared_memory.SharedMemory(name=f"out-{self.name}", create=False)
        self.out_np_shm = np.ndarray((20,6), dtype=np.float32, buffer=self.out_shm.buf)

    def detect(self, tensor_input, threshold=.4):
        detections = []

        # copy input to shared memory
        self.np_shm[:] = tensor_input[:]
        self.event.clear()
        self.detection_queue.put(self.name)
        result = self.event.wait(timeout=10.0)

        # if it timed out
        if result is None:
            return detections

        for d in self.out_np_shm:
            if d[1] < threshold:
                break
            detections.append((
                self.labels[int(d[0])],
                float(d[1]),
                (d[2], d[3], d[4], d[5])
            ))
        self.fps.update()
        return detections

    def cleanup(self):
        self.shm.unlink()
        self.out_shm.unlink()