import torch
from models import build_network
from torch.utils.data import Dataset
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
from tqdm import tqdm
from utils.loss import getLanMask
from pathlib import Path
from easydict import EasyDict
from utils.config import cfg_from_yaml_file
from transformers import AutoTokenizer
from torch import Tensor

def load_model(cfg, checkpoint_pth='../BriVL-pretrain-model/BriVL-1.0-5500w.pth'):
    model = build_network(cfg.MODEL)
    model_component = torch.load(checkpoint_pth, map_location=torch.device('cpu'))
    model.learnable.load_state_dict(model_component['learnable'])  ####### only save learnable
    model = torch.nn.DataParallel(model, device_ids=[cfg.DEVICE.GPU])
    if torch.cuda.is_available():
        model = model.cuda()
    model.eval()
    return model

def encode_images(cfg, images_pth, model):
    dataset = ImageDataset(images_pth, cfg)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        # num_workers=cfg.DATASET.WORKERS,
        pin_memory=True,
        drop_last=False
    )
    with torch.no_grad():
        np_text, np_img = None, None
        for idx, batch in enumerate(tqdm(dataloader)):
            # data
            imgs = batch[0]
            img_lens = batch[1].view(-1)
            image_boxs = batch[2]
            imgMask = getLanMask(img_lens, cfg.MODEL.MAX_IMG_LEN)
            if torch.cuda.is_available():
                imgMask = imgMask.cuda()
                imgs = imgs.cuda()
                image_boxs = image_boxs.cuda()  # <BSZ, 36, 4>

            # text_lens = text_lens.cuda() ############
            img = model(imgs, None, imgMask, None, None, image_boxs, is_training=False, schema='image')

            if np_img is None:
                np_img = img.cpu().numpy()  # <bsz, featdim>

            else:
                np_img = np.concatenate((np_img, img.cpu().numpy()), axis=0)

    return np_img

def encode_texts(cfg, texts, model):
    dataset = TextDataset(texts, cfg=cfg)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        # num_workers=cfg.DATASET.WORKERS,
        pin_memory=True,
        drop_last=False
    )
    with torch.no_grad():
        np_text = None
        for idx, batch in enumerate(tqdm(dataloader)):
            texts = batch[0]
            text_lens = batch[1]
            textMask = getLanMask(text_lens, max_len=cfg.MODEL.MAX_TEXT_LEN)
            if torch.cuda.is_available():
                textMask = textMask.cuda()
                texts = texts.cuda()
                text_lens = text_lens.cuda()
            text = model(None, texts, None, textMask, text_lens, None, is_training=False, schema='text')

            if np_text is None:
                np_text = text.cpu().numpy()
            else:
                np_text = np.concatenate((np_text, text.cpu().numpy()), axis=0)
    return np_text

class TextDataset(Dataset):
    def __init__(self, texts, cfg):
        self.texts = texts
        self.max_length = cfg.MODEL.MAX_TEXT_LEN
        self.text_transform = AutoTokenizer.from_pretrained(cfg.MODEL.ENCODER)

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

    def __getitem__(self, index):
        text = self.texts[index]

        text_info = self.text_transform(text, padding='max_length', truncation=True,
                                        max_length=self.max_length, return_tensors='pt')
        text = text_info.input_ids.reshape(-1)
        text_len = torch.sum(text_info.attention_mask)
        return text, text_len

class ImageDataset(Dataset):
    def __init__(self, images_pth, cfg):
        self.images_pth = images_pth
        self.new_size = cfg.MODEL.IMG_SIZE
        self.cfg = cfg
        mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
        normalize = transforms.Normalize(mean=mean, std=std)
        self.visual_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((self.new_size, self.new_size)),
            normalize])

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

    def __getitem__(self, index):
        # new_size = self.cfg.MODEL.IMG_SIZE
        img_path = self.images_pth[index]
        image = Image.open(img_path).convert('RGB')
        # width, height = image.size

        img_box_s = [torch.from_numpy(np.array([0, 0, self.new_size, self.new_size]).astype(np.float32))]

        valid_len = len(img_box_s)
        img_len = torch.full((1,), valid_len, dtype=torch.long)
        if valid_len < self.cfg.MODEL.MAX_IMG_LEN:
            for i in range(self.cfg.MODEL.MAX_IMG_LEN - valid_len):
                img_box_s.append(torch.from_numpy(np.array([0, 0, 0, 0]).astype(np.float32)))
        image_boxs = torch.stack(img_box_s, 0)

        image = self.visual_transform(image)

        return image, img_len, image_boxs

def cos_sim(a: Tensor, b: Tensor):
    """
    Computes the cosine similarity cos_sim(a[i], b[j]) for all i and j.
    :return: Matrix with res[i][j]  = cos_sim(a[i], b[j])
    """
    if not isinstance(a, torch.Tensor):
        a = torch.tensor(a)

    if not isinstance(b, torch.Tensor):
        b = torch.tensor(b)

    if len(a.shape) == 1:
        a = a.unsqueeze(0)

    if len(b.shape) == 1:
        b = b.unsqueeze(0)

    a_norm = torch.nn.functional.normalize(a, p=2, dim=1)
    b_norm = torch.nn.functional.normalize(b, p=2, dim=1)
    return torch.mm(a_norm, b_norm.transpose(0, 1))

if __name__ == '__main__':

    cfg = EasyDict()
    cfg.ROOT_DIR = (Path(__file__).resolve().parent / './').resolve()
    cfg_from_yaml_file('./cfg/test_xyb.yml', cfg)

    if torch.cuda.is_available():
        torch.cuda.set_device(cfg.DEVICE.GPU)

    model = load_model(cfg, '../BriVL-pretrain-model/BriVL-1.0-5500w.pth')

    texts = ['日照香炉生紫烟', '遥看瀑布挂前川', '飞流直下三千尺', '疑是银河落九天',
             '千呼万唤始出来', '犹抱琵琶半遮面', '转轴拨弦三两声', '未成曲调先有情',
             '天生丽质难自弃', '一朝选在君王侧', '回眸一笑百媚生', '六宫粉黛无颜色'
             ]

    texts_emb = encode_texts(cfg, texts, model)
    print(texts_emb.shape)

    images_pth = ['data/imgs/img.png',
                  'data/imgs/img_1.png',
                  'data/imgs/img_2.png']

    images_emb = encode_images(cfg, images_pth, model)
    print(images_emb.shape)

    img = torch.from_numpy(images_emb)
    text = torch.from_numpy(texts_emb)

    if torch.cuda.is_available():
        img = img.cuda()
        text = text.cuda()

    # cos_sim_scores = cos_sim(img, text)

    cos_sim_scores = torch.zeros((len(images_pth), len(text)), dtype=torch.float32)  # .cuda()
    print('Pair-to-pair: calculating scores')
    for i in tqdm(range(len(images_pth))):  # row: image  col: text
        cos_sim_scores[i, :] = torch.sum(img[i] * text, -1)
        # print(cos_sim_scores[i, :])

    for path, scores in tqdm(zip(images_pth, cos_sim_scores)):
        print("Path:", path)
        print('Scores:', scores)
        sort = torch.argsort(scores, descending=True)
        print("Sort:", sort)
        print("Sorted texts:", [texts[i] for i in sort], '\n')


#!/usr/bin/env python
from flask import Flask

from flask import render_template, request
import requests
import pandas as pd
import numpy as np
import csv
import pickle
from collections import Counter
import nltk
import textstat as ts
import emoji
import regex
import re
from nltk.sentiment.vader import SentimentIntensityAnalyzer as SIA
import shap
from io import BytesIO
import base64
import urllib.parse
import matplotlib.pyplot as plt
import seaborn as sns
import xgboost

#Initialize app
# application = Flask(__name__, static_url_path='/static')
app = Flask(__name__, static_url_path='/static')

### Load the pickled objects for my model and other needed parts.
model = pickle.load(open("/home/ubuntu/canisaythat_aws/model/test_XGBoost_model.sav", 'rb'))
scaler = pickle.load(open("/home/ubuntu/canisaythat_aws/model/test_XGBoost_scaler.sav", 'rb'))
final_column_order = pickle.load(open("/home/ubuntu/canisaythat_aws/model/finalized_column_order2.sav", 'rb'))
source_averages = pickle.load(open("/home/ubuntu/canisaythat_aws/model/finalized_source_averages2.sav", 'rb'))
source_averages = source_averages.round(decimals = 2)
source_averages = source_averages.reset_index(drop=True)
source_averages.columns = final_column_order[1:55]

id_to_source = {0: 'Extremely Casual', 1: 'Company IM', 2: 'Workplace Casual', 3: 'Reports', 4: 'Dissertations'}
source_to_id = {'Extremely Casual': 0 ,'Company IM': 1, 'Workplace Casual': 2, 'Reports': 3, 'Dissertations': 4}

"""
Custom functions for feature engineering
"""
def textstat_stats(text):
    doc_length = len(text.split())
    flesch_ease = ts.flesch_reading_ease(text) #Flesch Reading Ease Score
    flesch_grade = ts.flesch_kincaid_grade(text) #Flesch-Kincaid Grade Level
    gfog = ts.gunning_fog(text) # FOG index, also indicates grade level
    auto_readability = ts.automated_readability_index(text) #approximates the grade level needed to comprehend the text.
    cl_index = ts.coleman_liau_index(text) #grade level of the text using the Coleman-Liau Formula.
    lw_formula = ts.linsear_write_formula(text) #grade level using the Linsear Write Formula.
    dcr_score = ts.dale_chall_readability_score(text) #uses a lookup table of the most commonly used 3000 English words
    syll_count = ts.syllable_count(text, lang='en_US')
    syll_count_scaled = syll_count / doc_length
    lex_count = ts.lexicon_count(text, removepunct=True)
    lex_count_scaled = lex_count / doc_length
    idx = ['flesch_ease', 'flesch_grade','gfog',
           'auto_readability','cl_index','lw_formula',
           'dcr_score',
           'syll_count', 'lex_count']
    return pd.Series([flesch_ease, flesch_grade, gfog,
                      auto_readability, cl_index, lw_formula,
                      dcr_score,
                      syll_count_scaled, lex_count_scaled], index = idx)

def emoji_counter(text):
    emoji_list = []
    data = regex.findall(r'\X', text)
    for word in data:
        if any(char in emoji.UNICODE_EMOJI for char in word):
            emoji_list.append(word)
    return emoji_list

def scream_counter(text):
    crazy_confused = len(re.findall(r'!\?|\?!', text))
    ahhhh = len(re.findall(r'!!', text))
    huhhh = len(re.findall(r'\?\?', text))
    screams = crazy_confused + ahhhh + huhhh
    return screams

Google_Curses = pd.read_csv("/home/ubuntu/canisaythat_aws/model/custom_curse_words.txt", header = None)
bad_words = Google_Curses[0].tolist()

def swear_counter(text): #returns number of curses in text
    swear_jar = []
    what_curses = []
    for word in bad_words:
        curses = text.count(word)
        swear_jar.append(curses)
        if curses > 0:
            what_curses.append(word)
    return sum(swear_jar), what_curses

contractions = [
  r'n\'t',
  r'I\'m',
  r'(\w+)\'ll',
  r'(\w+)n\'t',
  r'(\w+)\'ve',
  r'(\w+)\'s',
  r'(\w+)\'re',
  r'(\w+)\'d',
]

def contraction_counter(text):
    doc_length = len(text.split())
    abbreviations = []
    for abbrev in contractions:
        num_abbrevs = len(re.findall(abbrev, text))
        abbreviations.append(num_abbrevs)
    return sum(abbreviations) / doc_length

"""
Process the user text, to generate features used in model
"""
def process_user_text(user_text, goal_category):

    ## put user input text string into a DataFrame
    clean_data = pd.DataFrame(user_text, columns = ["text"])
    clean_data["source"] = goal_category
    # clean_data["subreddit"] = "placeholder"

    # remove hyperlinks & bullets (so they don't get counted as emoji)
    clean_data["text"] = clean_data["text"].str.replace(r'http\S*\s', ' ')
    clean_data["text"] = clean_data["text"].str.replace(r'http\S*(\n|\)|$)', ' ')
    clean_data["text"] = clean_data["text"].str.replace(r'', ' ')

    ## Starting with textstat
    textstat_results = pd.DataFrame(columns = ['flesch_ease', 'flesch_grade','gfog',
               'auto_readability','cl_index','lw_formula','dcr_score',
               'syll_count', 'lex_count'])

    for i in clean_data["text"]:
        results = textstat_stats(str(i)) #textstat needs a string
        textstat_results = textstat_results.append(results, ignore_index=True) #so that index is continuous

    # Resetting indices here may be unneccesary
    textstat_results = textstat_results.reset_index(drop=True)

    combined_data = pd.concat([clean_data, textstat_results], axis = 1)

    ## Moving on to NLTK part-of-speech tagging
    combined_data_wordtokens = []
    for document in combined_data["text"]:
        tokens = nltk.word_tokenize(str(document))
        combined_data_wordtokens.append(tokens)

    combined_data_wordpos = []
    for document in combined_data_wordtokens:
        pos = nltk.pos_tag(document) #default is Penn Treebank tagset
        combined_data_wordpos.append(pos)

    pos_counts = []

    pos_keys = ['CC', 'CD','DT','EX','FW','IN', 'JJ','JJR','JJS','LS','MD','NN','NNS',
               'NNP','NNPS','PDT','POS','PRP','PRP$','RB','RBR' ,'RBS','RP', 'SYM','TO',
               'UH','VB', 'VBD', 'VBG','VBN','VBP','VBZ','WDT','WP','WP$','WRB',
               '#', '$', "''", '(', ')', ',', '.', ':','``']

    for document in combined_data_wordpos:
        doc_length = len(document)
        mini_dict = Counter([pos for word,pos in document])
        for pos in pos_keys:
           if pos not in mini_dict:
               mini_dict[pos] = 0
        scaled_dict = {k: v / doc_length for k, v in mini_dict.items()}
        pos_counts.append(scaled_dict)

    # for document in combined_data_wordpos:
    #     doc_length = len(document)
    #     mini_dict = Counter([pos for word,pos in document])
    #     scaled_dict = {k: v / doc_length for k, v in mini_dict.items()}
    #     pos_counts.append(scaled_dict)

    pos_df = pd.DataFrame(pos_counts)
    pos_df = pos_df.fillna(0)

    combined_data = pd.concat([combined_data, pos_df], axis = 1)

    ## Add sentiment intensity
    sia = SIA()
    results = []

    for line in combined_data["text"]: # this returns a list of dicts
        pol_score = sia.polarity_scores(line)
        pol_score['text'] = line
        results.append(pol_score)

    sia_neg = []
    sia_pos = []
    sia_neu = []
    sia_comp = []
    for document in results:
        neg = document['neg']
        pos = document['pos']
        neu = document['neu']
        comp = document['compound']
        sia_neg.append(neg)
        sia_pos.append(pos)
        sia_neu.append(neu)
        sia_comp.append(comp)

    combined_data["SIA_neg"] = sia_neg
    combined_data["SIA_pos"] = sia_pos
    combined_data["SIA_neu"] = sia_neu
    combined_data["SIA_com"] = sia_comp

    ## Now for the custom features
    #  cursing
    swear_jar = []
    what_curses = []
    for row in combined_data["text"]:
        num_curses, which_curses = swear_counter(row)
        swear_jar.append(num_curses)
        what_curses.append(which_curses)
    combined_data["num_curses"] = swear_jar

    # emoji use
    emoji_counts = []
    for row in combined_data["text"]:
        emoji_num = len(emoji_counter(str(row)))
        emoji_counts.append(emoji_num)
    combined_data["num_emoji"] = emoji_counts

    # !?,  ?!, ??, and !!
    internet_yelling = []
    for row in combined_data["text"]:
        screams = scream_counter(str(row))
        internet_yelling.append(screams)
    combined_data["yell_count"] = internet_yelling

    # Combine the punctuation columns to get a total measure
    punct_columns = ['#', '$', "''", '(', ')', ',', '.', ':','``']
    combined_data['total_punctuation'] = combined_data.loc[:, punct_columns].sum(axis=1)

    # Combine parentheses
    parentheses = ['(',')']
    combined_data['parentheses'] = combined_data.loc[:, parentheses].sum(axis=1)

    # Encoding of quotation marks varies by source, need to combine them.
    quotes = ["''",'``', "'"]
    combined_data['quotes'] = combined_data.loc[:, quotes].sum(axis=1)

    # Count up contractions, scale by doc length
    num_abbreviations = []
    for row in combined_data["text"]:
        num_abbrevs = contraction_counter(str(row))
        num_abbreviations.append(num_abbrevs)
    combined_data["contractions"] = num_abbreviations

    # Google_Curses = pd.read_csv("./model/custom_curse_words.txt", header = None)
    # bad_words = Google_Curses[0].tolist()

    # any_bad = []
    # for row in combined_data["text"]:
    #     if any(str(word) in str(row) for word in bad_words):
    #         any_bad.append(1)
    #     else: any_bad.append(0)

    # combined_data["Google_curses"] = any_bad
    # combined_data["Google_curses"].value_counts()

    # emoji_counts = []
    # for row in combined_data["text"]:
    #     emoji_num = len(emoji_counter(str(row)))
    #     emoji_counts.append(emoji_num)

    # combined_data["Num_emoji"] = emoji_counts
    # combined_data["Num_emoji"].value_counts()

    # internet_yelling = []
    # for row in combined_data["text"]:
    #     screams = scream_counter(str(row))
    #     internet_yelling.append(screams)

    # combined_data["Yell_count"] = internet_yelling

    ## Make sure user data is in the same order as the model.
    features_to_drop = ['gfog','``','auto_readability','lex_count', '#',
                        'LS', 'SYM','WP$', '``', "''",'(', ')']

    combined_data = combined_data.drop(features_to_drop, axis = 1)

    # this may be extremely important. starts with text, ends with source
    combined_data = combined_data[final_column_order]

    return combined_data

def get_professionalism_score(user_df):
    combined_data = user_df

    # split data into X and y
    userX = combined_data.drop(["text", "source"], axis = 1)
    userY = combined_data['source'] # "source" is the goal category

    col_names = userX.columns # will user later to reattach names

    # scale the data using scaler trained on original corpus
    userX = scaler.transform(userX)

    # make predictions for test data
    y_probs = model.predict_proba(userX)

    top_two_classes = y_probs[0].argsort()[-2:][::-1]
    first_class_prob = y_probs[0][top_two_classes[0]]
    first_class_prob  = round(first_class_prob * 100)
    second_class_prob = y_probs[0][top_two_classes[1]]
    second_class_prob  = round(second_class_prob * 100)

    #attach column names to scaled features for user
    userX_df = pd.DataFrame(userX)
    userX_df.columns = col_names
    userX_df = userX_df.round(decimals = 2)
    userX_df = userX_df.reset_index(drop=True)

    # class_diff = first_class_prob - second_class_prob

    # if class_diff > 0.2: # if the choice was clear, go with highest prob class
    #     user_prof_score = top_two_classes[0]
    # else:  # if choice was close, go with a value in between
    #     user_prof_score = min(top_two_classes[0], top_two_classes[1]) + 0.5

    # return user_prof_score
    return top_two_classes, first_class_prob, second_class_prob, userX_df, y_probs

def make_classification_plot(y_probs, top_class_numeric):
    y_probs2 = [i for i in y_probs[0]]
    df = pd.DataFrame(y_probs2, columns = ["probs"])
    df["source"] = source_to_id
    df["source2"] = pd.Categorical(df["source"])
    my_range=range(1,len(df.index)+1)

    # Create a color if the source is the highest probability class
    my_color=np.where(df ['source2']==id_to_source[top_class_numeric], 'orange', 'skyblue')
    my_size=np.where(df ['source2']==id_to_source[top_class_numeric], 70, 30)

    # set up the thing that will hold my figure
    img = BytesIO()

    # image size
    fig, ax = plt.subplots(figsize=[6,2], dpi = 300)
    plt.tight_layout()

    # The vertival plot is made using the hline function
    plt.hlines(y=my_range, xmin=0, xmax=df['probs']*100, color=my_color, alpha=0.4)
    plt.scatter(df['probs']*100, my_range, color=my_color, s=my_size, alpha=1)

    # Add title and exis names
    plt.yticks(my_range, df['source'])
    # plt.title(, loc='left')
    plt.xlabel("Percent match to each category")
    plt.ylabel(None)

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)

    # add the plot data to the img thing
    plt.savefig(img, format='png', bbox_inches = 'tight')
    plt.close()
    img.seek(0)
    img_png_target = img.getvalue()
    plot_url = base64.b64encode(img_png_target)

    return plot_url
    # plt.savefig('./insight2019/flask_app/my_flask/model/test_plot.png', bbox_inches='tight')

def get_SHAP_results(user_df, user_category, goal_category):

    #set up dictionary of feature explanations.
    with open("/home/ubuntu/canisaythat_aws/model/column_keys.txt") as fin:
         rows = ( line.strip().split('\t') for line in fin )
         column_dict = { row[0]:row[1:] for row in rows }

    #starting with output from process_user_text()
    # X_train_labeled = pd.DataFrame(user_df, columns = user_df.columns[3:])
    user_df = user_df.drop(["text", "source"], axis = 1) # now we have only feature columns
    X_train_labeled = user_df

    #set up the SHAP explainer for the XGBoost model and calc SHAP scores
    explainer = shap.TreeExplainer(model)
    shap_values = explainer.shap_values(X_train_labeled) # list of lists, outer lists
    # are the categories, inner list is feature contributions to that category

    # re-arrange SHAP values so they're easy to give back to user
    # source0_shap=pd.DataFrame(data=shap_values[0], columns=user_df.columns[3:], index=[0])
    # source1_shap=pd.DataFrame(data=shap_values[1], columns=user_df.columns[3:], index=[1])
    # source2_shap=pd.DataFrame(data=shap_values[2], columns=user_df.columns[3:], index=[2])
    # source3_shap=pd.DataFrame(data=shap_values[3], columns=user_df.columns[3:], index=[3])
    # source4_shap=pd.DataFrame(data=shap_values[4], columns=user_df.columns[3:], index=[4])
    source0_shap=pd.DataFrame(data=shap_values[0], columns=user_df.columns, index=[0])
    source1_shap=pd.DataFrame(data=shap_values[1], columns=user_df.columns, index=[1])
    source2_shap=pd.DataFrame(data=shap_values[2], columns=user_df.columns, index=[2])
    source3_shap=pd.DataFrame(data=shap_values[3], columns=user_df.columns, index=[3])
    source4_shap=pd.DataFrame(data=shap_values[4], columns=user_df.columns, index=[4])

    user_reasoning=pd.concat([source0_shap,
                              source1_shap,
                              source2_shap,
                              source3_shap,
                              source4_shap])
    user_reasoning = user_reasoning.transpose()

    # make sure to translate the written sources into numeric
    goal_numeric = source_to_id[goal_category]
#    user_cat_string= id_to_source[user_category] #testing

    #Why did user text get it's particular score? (top 3 features)
#    user_category_reasoning = user_reasoning.sort_values(by=[user_cat_string], ascending=False)
    user_category_reasoning = user_reasoning.sort_values(by=user_category, ascending=False)
    user_category_reasoning_labels = user_category_reasoning.index.values[0:3]
    user_category_reasoning_labels2 = [column_dict[feature] for feature in user_category_reasoning_labels]
    user_category_reasoning_labels2 = [item for sublist in user_category_reasoning_labels2 for item in sublist]
    user_category_reasoning_avgs = source_averages.loc[user_category, user_category_reasoning_labels]

    # What features diminished their odds of being labeled as their goal?  (top 3 features)
    user_goal_improvement = user_reasoning.sort_values(by=[goal_numeric], ascending=True)
    user_goal_improvement_labels = user_goal_improvement.index.values[0:3]
    user_goal_improvement_labels2 = [column_dict[feature] for feature in user_goal_improvement_labels]
    user_goal_improvement_labels2 = [item for sublist in user_goal_improvement_labels2 for item in sublist]
    user_goal_improvement_avgs = source_averages.loc[goal_numeric, user_goal_improvement_labels]

    # What features improved their odds of being labeled as their goal?  (top 3 features)
    user_goal_encouragement = user_reasoning.sort_values(by=[goal_numeric], ascending=False)
    user_goal_encouragement_labels = user_goal_encouragement.index.values[0:3]
    user_goal_encouragement_labels2 = [column_dict[feature] for feature in user_goal_encouragement_labels]
    user_goal_encouragement_labels2 = [item for sublist in user_goal_encouragement_labels2 for item in sublist]
    user_goal_encouragement_avgs = source_averages.loc[goal_numeric, user_goal_encouragement_labels]

    return user_category_reasoning_labels, user_category_reasoning_labels2, user_category_reasoning_avgs, user_goal_improvement_labels, user_goal_improvement_labels2, user_goal_improvement_avgs, user_goal_encouragement_labels, user_goal_encouragement_labels2, user_goal_encouragement_avgs

def make_feedback_lollipop(improvement_feedback2, imp_cat_val, improvement_avgs, encouragement_feedback2, enc_cat_val, encouragement_avgs):
    lollipop_red = pd.DataFrame([improvement_feedback2, imp_cat_val, improvement_avgs])
    lollipop_red = lollipop_red.transpose()
    lollipop_red.columns = ["feature_name", "user_score", "category_average"]
    lollipop_red["purpose"] = "critique"

    lollipop_green  = pd.DataFrame([encouragement_feedback2, enc_cat_val, encouragement_avgs])
    lollipop_green = lollipop_green.transpose()
    lollipop_green.columns = ["feature_name", "user_score", "category_average"]
    lollipop_green["purpose"] = "encouragement"

    lollipops = pd.concat([lollipop_green, lollipop_red], axis = 0)
    lollipops["user_diff"] = lollipops["user_score"] - lollipops["category_average"]

    # add a threshold so that plots stay interpretable
    lollipops["user_diff"] = lollipops["user_diff"].clip(-2,2)

    # Create a color if the source is the highest probability class
    my_range=range(1,len(lollipops.index)+1)
    my_color=np.where(lollipops ['purpose']=="critique", 'purple', 'green')
    #my_size=np.where(lollipops ['source2']==id_to_source[top_class_numeric], 70, 30)

    # set up the thing that will hold my figure
    img = BytesIO()

    # The vertival plot is made using the hline function
    fig, ax = plt.subplots(figsize=(6, 2.4), dpi = 300)
    plt.tight_layout()

    plt.vlines(x=0, ymin = min(my_range), ymax=max(my_range), linestyles="dashed", color = "silver", zorder=0)
    plt.hlines(y=my_range, xmin=0, xmax=lollipops['user_diff'], color=my_color, alpha=0.4)
    plt.scatter(lollipops['user_diff'],
                my_range,
                color=my_color,
                alpha=1)

    # Add title and exis names
    plt.yticks(my_range, lollipops['feature_name'])
    # plt.title(, loc='left')
    plt.xlabel("Difference from goal category average")
    plt.ylabel(None)
    plt.xlim(-2,2)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    # ax.spines['bottom'].set_visible(False)
    # ax.set_xticks([-2,0,2])
    plt.xticks([-2,0,2], ('Too low', 'Just right', 'Too high'))
    # ax.get_xaxis().set_visible(False)

    # add the plot data to the img thing
    plt.savefig(img, format='png', bbox_inches = 'tight')
    plt.close()
    img.seek(0)
    img_png_target = img.getvalue()
    plot_url = base64.b64encode(img_png_target)

    return plot_url

#Standard home page. 'index.html' is the file in your templates that has the CSS and HTML for your app
# @app.route('/', methods=['GET', 'POST'])
# def index():
#     return render_template('index.html')

# @app.route('/input', methods=['GET', 'POST'])

@app.route('/', methods=['GET', 'POST'])
def get_inputs():
  sources = list(id_to_source.values())
  return render_template("inputs.html", goal_category = sources)

@app.route('/output', methods=['GET', 'POST'])
def get_outputs():

    goal_category = str(request.form['goal_category'])
    user_text = str(request.form['user_text'])
    # user_text = "This is a placeholder sentence."

    # read data (should be a string) and make it a list of length 1.
    # user_text1 = [i for i in user_text]
    user_text1 = [user_text]

    # process the text, get back a dataframe
    user_df = process_user_text(user_text1, goal_category) # Returns a DataFrame with 1 row

    # run the model function
    top_two_classes, first_class_prob, second_class_prob, userX, y_probs = get_professionalism_score(user_df)
    top_class = id_to_source[top_two_classes[0]]
    second_class = id_to_source[top_two_classes[1]]

    # make a plot to show probability of classification in each category
    class_plot = make_classification_plot(y_probs, top_two_classes[0])

    # Run SHAP, get 3 lists of 3 features each
    # For those features, look up their average for either the goal or user category
    category_reasoning, category_reasoning2, category_avgs, improvement_feedback, improvement_feedback2, improvement_avgs, encouragement_feedback, encouragement_feedback2, encouragement_avgs = get_SHAP_results(user_df, top_two_classes[0], goal_category)

    # Get user values for all 9 features I want to report
    user_cat_val = userX.loc[0, category_reasoning]
    imp_cat_val = userX.loc[0, improvement_feedback]
    enc_cat_val = userX.loc[0, encouragement_feedback]

    # make a results plot that summarizes the good and bad
    feedback_plot = make_feedback_lollipop(improvement_feedback2, imp_cat_val, improvement_avgs, encouragement_feedback2, enc_cat_val, encouragement_avgs)

    return render_template(
      "outputs.html",
      class_plot_url = urllib.parse.quote(class_plot),
      feedback_plot_url = urllib.parse.quote(feedback_plot),
      values = {
          'user_text': user_text if user_text else "No input.",
          'goal_category': goal_category if goal_category else "No input.",
          # 'class_plot': class_plot,
          'top_class': top_class,
          'second_class': second_class,
          'top_class_prob': first_class_prob,
          'second_class_prob': second_class_prob,
          'category_reasoning': category_reasoning2,
          'user_cat_val': user_cat_val,
          'category_avgs': category_avgs,
          'improvement_feedback': improvement_feedback2,
          'imp_cat_val': imp_cat_val,
          'improvement_avgs': improvement_avgs,
          'encouragement_feedback': encouragement_feedback2,
          'enc_cat_val': enc_cat_val,
          'encouragement_avgs': encouragement_avgs
      })

# 'prof_score': prof_score if prof_score else "No result."

# @app.route('/input')
# def get_inputs():
#     print(request.args)

#     user_text = request.args.get('user_text')
#     # user_text = "This is a placeholder sentence."
#     prof_score = "placeholder"

#     # read data (should be a string) and make it a list.
#     user_text1 = [i for i in user_text]

#     # run the model function
#     prof_score = get_professionalism_score(user_text1)

#     return render_template(
#       "inputs.html",
#        title = 'Home',
#        values = {
#           'user_text': user_text if user_text else "No input.",
#           'prof_score': prof_score if prof_score else "No result."
#       },
#     )

if __name__ == '__main__':
    # this runs your app locally:
    # application.run(host='0.0.0.0', port=8080, debug=True)
    app.run(host='0.0.0.0',debug=True)


#coding=utf-8

import os
from torch import nn
import warnings
import pandas as pd
import numpy as np
from torch.utils.data import Dataset
from sklearn.preprocessing import MultiLabelBinarizer
from utils.feature_utils import read_all_features
import random
import torch
from torch.utils.data import DataLoader

import argparse

# 1. set random seed
random.seed(2050)
np.random.seed(2050)
torch.manual_seed(2050)
torch.cuda.manual_seed_all(2050)

torch.backends.cudnn.benchmark = True
warnings.filterwarnings('ignore')

parser = argparse.ArgumentParser(description='Finetune models Test')

parser.add_argument('--gpuid', '-g', default='0', metavar='N',
                    help='gpu id used for training (default: 0)')
parser.add_argument('--tag', '-t', default='', help='tag used for save')

class BaseConfig():
    def __init__(self,tag=''):
        self.features_type='inceptionv3_inceptionv4_xception_800_offi_hpa'
        self.num_classes = 28
        self.batch_size = 512
        self.sample_submit_file = "./data/sample_submission.csv"
        base_dir= './models/MLP/'
        self.test_file = './data/sample_submission_featureindex_label.csv'
        self.features_test = './features/inceptionv3_inceptionv4_xception_800_10fold/test_offi_hpa_features.h5'
        self.feature_dim = 3840
        print('[features_test]:', self.features_test)
        if tag != '':
            base_dir = base_dir+'/'+tag
        self.weights = base_dir + "/models"
        self.submit = base_dir + "/submit/"
        self.logs_dir = base_dir + '/logs'
        # self.best_models = base_dir+'/bestmodel'
        self.submit_file = os.path.join(self.submit, '{}_submission.csv'.format(self.features_type))

        self.run_type = 'train'
        mkdir_with_check(base_dir)
        mkdir_with_check(self.weights)
        mkdir_with_check(self.submit)
        mkdir_with_check(self.logs_dir)
        # mkdir_with_check(self.best_models)

def mkdir_with_check(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
        try:
            os.system('chmod 777 {}'.format(dir_path))
        except:
            pass

class HumanDataset_MLP(Dataset):
    def __init__(self, images_df, mode="test", config=''):
        self.config = config
        self.images_df = images_df.copy()
        self.mlb = MultiLabelBinarizer(classes=np.arange(0, self.config.num_classes))
        self.mlb.fit(np.arange(0, self.config.num_classes))
        self.mode = mode
        if mode=='train':
            self.features = read_all_features(self.config.features_train)
        else:
            self.features = read_all_features(self.config.features_test)

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

    def __getitem__(self, index):
        if not isinstance(index, int):
            index = index.item()
        feature_index = int(self.images_df.iloc[index].Id)
        X = self.features[feature_index]
        if not self.mode == "test":
            labels = np.array(list(map(int, self.images_df.iloc[index].Target.split(' '))))
            y = np.eye(self.config.num_classes, dtype=np.float)[labels].sum(axis=0)
        else:
            y = str(self.images_df.iloc[index].Id)
        X = torch.from_numpy(X)
        return X.float(), y

class MLP(nn.Module):
    def __init__(self, config):
        super(MLP, self).__init__()
        self.fc1 = nn.Sequential(
            nn.BatchNorm1d(config.feature_dim),
            nn.Dropout(0.5),
            nn.Linear(config.feature_dim, 128))
        self.fc2 = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, config.num_classes),
        )

    def forward(self, input):

        x_feature = self.fc1(input)
        x_class = self.fc2(x_feature)

        return x_class, x_feature
# 3. test model on public dataset and save the probability matrix
def test(test_loader, model, fold,write_result=True):
    sample_submission_df = pd.read_csv(config.sample_submit_file)
    filenames, labels, submissions = [], [], []
    model.cuda()
    model.eval()
    submit_results = []
    result_score = np.zeros((len(test_loader), config.num_classes))
    for i, (input, filepath) in enumerate(test_loader):
        # 3.2 change everything to cuda and get only basename
        filepath = [os.path.basename(x) for x in filepath]
        with torch.no_grad():
            image_var = input.cuda(non_blocking=True)
            y_pred, sub_features = model(image_var)
            label = y_pred.sigmoid().cpu().data.numpy()
            result_score[i] = label[0]
            # print(label > 0.5)
            labels.append(label > 0.15)
            filenames.append(filepath)

    if write_result:
        for row in np.concatenate(labels):
            subrow = ' '.join(list([str(i) for i in np.nonzero(row)[0]]))
            submissions.append(subrow)
        sample_submission_df['Predicted'] = submissions
        sample_submission_df.to_csv(config.submit_file[:-4] + '_fold{}.csv'.format(fold), index=None)
        np.save(config.submit_file[:-4] + '_fold{}_score.npy'.format(fold), result_score)
        print('result has saved to ' + config.submit_file[:-4] + '_fold{}.csv'.format(fold))

# 4. main function
def predict(fold):

    # 4.1 mkdirs
    if not os.path.exists(config.submit):
        os.makedirs(config.submit)

    model = MLP(config)
    model.cuda()

    test_files = pd.read_csv(config.test_file)
    test_gen = HumanDataset_MLP(test_files, mode="test", config=config)
    test_loader = DataLoader(test_gen, 1, shuffle=False, pin_memory=True, num_workers=0)
    model_path=os.path.join(config.weights,'{}_fold_{}_model_best_loss.pth.tar'.format(config.features_type,fold))

    print('load model from:'+model_path)

    best_model = torch.load(model_path)
    model.load_state_dict(best_model["state_dict"])
    test(test_loader, model, fold)

def main():
    for i in xrange(10):
        predict(i)

if __name__ == "__main__":
    args = parser.parse_args()
    gpu_id = args.gpuid
    config = BaseConfig(tag=args.tag)
    # model_input=args.model_path
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    main()


import os
import numpy as np
import tensorflow as tf
from tensorflow import keras

from actor_critic import Actor, Critic
from replay_buffer import ReplayBuffer

REPLAY_BUFFER_SIZE = 100000
REPLAY_START_SIZE = 10000
BATCH_SIZE = 128
GAMMA = 0.99

class DDPG:
    def __init__(
        self,
        env,
        obs_dim,
        action_dim,
        action_gain=1,
        actor_learning_rate=0.0001,
        critic_learning_rate=0.001,
        gamma=0.99,
        tau=0.001,
    ):
        self.environment = env
        self.obs_dim = obs_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.tau = tau
        self.time_step = 0

        # make main networks
        self.actor = Actor(obs_dim, action_dim, action_gain, actor_learning_rate)
        self.critic = Critic(obs_dim, action_dim, critic_learning_rate)

        # make target networks
        self.target_actor = Actor(obs_dim, action_dim, action_gain)
        self.target_actor.model.set_weights(self.actor.model.get_weights())
        self.target_critic = Critic(obs_dim, action_dim)
        self.target_critic.model.set_weights(self.critic.model.get_weights())

        # set replay buffer
        self.replay_buffer = ReplayBuffer(REPLAY_BUFFER_SIZE)

    def action(self, obs):
        return self.actor.act(obs)[0]

    def train(self):
        batch = self.replay_buffer.sample_batch(BATCH_SIZE)
        self.update_networks(batch)

    def perceive(self, state, action, reward, next_state, done):
        # Store transition (s_t,a_t,r_t,s_{t+1}) in replay buffer
        self.replay_buffer.add([state, action, reward, done, next_state])
        if self.replay_buffer.size() == REPLAY_START_SIZE:
            print('\n---------------Start training---------------')
        # Store transitions to replay start size then start training
        if self.replay_buffer.size() > REPLAY_START_SIZE:
            self.time_step += 1
            self.train()

        if self.time_step % 10000 == 0 and self.time_step > 0:
            self.save_model("./results")

        return self.time_step

    @tf.function
    def update_networks(self, batch):
        """ runs all updates from provided training batch """
        s, a, r, t, next_s = (tf.cast(i, tf.float32) for i in batch)
        self.update_critic(s, a, r, next_s)
        self.update_actor(s, a, r, next_s)
        self.update_target(self.actor.model, self.target_actor.model)
        self.update_target(self.critic.model, self.target_critic.model)

    @tf.function
    def update_critic(self, s, a, r, next_s):
        """ minimize td-loss from target """
        # td estimate based on targets' behavior
        target_future_actions = self.target_actor.act(next_s)
        target_future_qs = self.target_critic.estimate_q(next_s, target_future_actions)
        target_current_qs = r + self.gamma * target_future_qs

        # update main critic
        main_current_qs = self.critic.model([s, a])
        loss = keras.losses.mse(target_current_qs, main_current_qs)

        model_vars = self.critic.model.trainable_variables

        dloss_dcrit = tf.gradients(loss, model_vars)
        self.critic.optimizer.apply_gradients(zip(dloss_dcrit, model_vars))

    @tf.function
    def update_actor(self, s, a, r, next_s):
        """ dq_dtheta = dq_da * da_dtheta"""
        # first, finding dq_da
        proposed_a = self.actor.model(s)
        q = self.critic.model([s, proposed_a])
        dq_da = tf.gradients(q, proposed_a)[0]

        # second, finding dq_da * da_dtheta
        model_vars = self.actor.model.trainable_variables
        dq_dtheta = tf.gradients(proposed_a, model_vars, grad_ys=-dq_da)

        # updating the model
        self.actor.optimizer.apply_gradients(zip(dq_dtheta, model_vars))

    @tf.function
    def update_target(self, main_model, target_model):
        """ target = tau*main + (1-tau)*target """
        for model_weight, target_weight in zip(
            main_model.weights, target_model.weights
        ):
            target_weight.assign(
                self.tau * model_weight + (1 - self.tau) * target_weight
            )

    def save_model(self, save_dir):
        """ saves the main and target networks"""
        self.actor.model.save_weights(os.path.join(save_dir, f"actor_{self.time_step}"))
        self.critic.model.save_weights(os.path.join(save_dir, f"critic_{self.time_step}"))
        self.target_actor.model.save_weights(os.path.join(save_dir, f"target_actor_{self.time_step}"))
        self.target_critic.model.save_weights(os.path.join(save_dir, f"target_critic_{self.time_step}"))


#!/usr/bin/env python3
""" S3FD Face detection plugin
https://arxiv.org/abs/1708.05237

Adapted from S3FD Port in FAN:
https://github.com/1adrianb/face-alignment
"""

from scipy.special import logsumexp

import numpy as np

from lib.multithreading import MultiThread
from ._base import BoundingBox, Detector, logger

class Detect(Detector):
    """ S3FD detector for face recognition """
    def __init__(self, **kwargs):
        git_model_id = 3
        model_filename = "s3fd_v1.pb"
        super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs)
        self.name = "s3fd"
        self.target = (640, 640)  # Uses approx 4 GB of VRAM
        self.vram = 4096
        self.min_vram = 1024  # Will run at this with warnings
        self.model = None

    def initialize(self, *args, **kwargs):
        """ Create the s3fd detector """
        try:
            super().initialize(*args, **kwargs)
            logger.info("Initializing S3FD Detector...")
            card_id, vram_free, vram_total = self.get_vram_free()
            if vram_free <= self.vram:
                tf_ratio = 1.0
            else:
                tf_ratio = self.vram / vram_total

            logger.verbose("Reserving %s%% of total VRAM per s3fd thread",
                           round(tf_ratio * 100, 2))

            confidence = self.config["confidence"] / 100
            self.model = S3fd(self.model_path, self.target, tf_ratio, card_id, confidence)

            if not self.model.is_gpu:
                alloc = 2048
                logger.warning("Using CPU")
            else:
                logger.debug("Using GPU")
                alloc = vram_free
            logger.debug("Allocated for Tensorflow: %sMB", alloc)

            if self.min_vram < alloc < self.vram:
                self.batch_size = 1
                logger.warning("You are running s3fd with %sMB VRAM. The model is optimized for "
                               "%sMB VRAM. Detection should still run but you may get "
                               "warnings/errors", int(alloc), self.vram)
            else:
                self.batch_size = int(alloc / self.vram)
            if self.batch_size < 1:
                raise ValueError("Insufficient VRAM available to continue "
                                 "({}MB)".format(int(alloc)))

            logger.verbose("Processing in %s threads", self.batch_size)

            self.init.set()
            logger.info("Initialized S3FD Detector.")
        except Exception as err:
            self.error.set()
            raise err

    def detect_faces(self, *args, **kwargs):
        """ Detect faces in Multiple Threads """
        super().detect_faces(*args, **kwargs)
        workers = MultiThread(target=self.detect_thread, thread_count=self.batch_size)
        workers.start()
        workers.join()
        sentinel = self.queues["in"].get()
        self.queues["out"].put(sentinel)
        logger.debug("Detecting Faces complete")

    def detect_thread(self):
        """ Detect faces in rgb image """
        logger.debug("Launching Detect")
        while True:
            item = self.get_item()
            if item == "EOF":
                break
            logger.trace("Detecting faces: '%s'", item["filename"])
            detect_image, scale = self.compile_detection_image(item["image"], is_square=True)
            for angle in self.rotation:
                current_image, rotmat = self.rotate_image(detect_image, angle)
                faces = self.model.detect_face(current_image)
                if angle != 0 and faces.any():
                    logger.verbose("found face(s) by rotating image %s degrees", angle)
                if faces.any():
                    break

            detected_faces = self.process_output(faces, rotmat, scale)
            item["detected_faces"] = detected_faces
            self.finalize(item)

        logger.debug("Thread Completed Detect")

    def process_output(self, faces, rotation_matrix, scale):
        """ Compile found faces for output """
        logger.trace("Processing Output: (faces: %s, rotation_matrix: %s)", faces, rotation_matrix)
        faces = [BoundingBox(face[0], face[1], face[2], face[3]) for face in faces]
        if isinstance(rotation_matrix, np.ndarray):
            faces = [self.rotate_rect(face, rotation_matrix)
                     for face in faces]
        detected = [BoundingBox(face.left / scale, face.top / scale,
                                face.right / scale, face.bottom / scale)
                    for face in faces]
        logger.trace("Processed Output: %s", detected)
        return detected

class S3fd():
    """ Tensorflow Network """
    def __init__(self, model_path, target_size, vram_ratio, card_id, confidence):
        logger.debug("Initializing: %s: (model_path: '%s', target_size: %s, vram_ratio: %s, "
                     "card_id: %s)",
                     self.__class__.__name__, model_path, target_size, vram_ratio, card_id)
        # Must import tensorflow inside the spawned process for Windows machines
        import tensorflow as tf
        self.is_gpu = False
        self.tf = tf  # pylint: disable=invalid-name
        self.model_path = model_path
        self.confidence = confidence
        self.graph = self.load_graph()
        self.input = self.graph.get_tensor_by_name("s3fd/input_1:0")
        self.output = self.get_outputs()
        self.session = self.set_session(target_size, vram_ratio, card_id)
        logger.debug("Initialized: %s", self.__class__.__name__)

    def load_graph(self):
        """ Load the tensorflow Model and weights """
        # pylint: disable=not-context-manager
        logger.verbose("Initializing S3FD Network model...")
        with self.tf.gfile.GFile(self.model_path, "rb") as gfile:
            graph_def = self.tf.GraphDef()
            graph_def.ParseFromString(gfile.read())
        fa_graph = self.tf.Graph()
        with fa_graph.as_default():
            self.tf.import_graph_def(graph_def, name="s3fd")
        return fa_graph

    def get_outputs(self):
        """ Return the output tensors """
        tensor_names = ["concat_31", "transpose_72", "transpose_75", "transpose_78",
                        "transpose_81", "transpose_84", "transpose_87", "transpose_90",
                        "transpose_93", "transpose_96", "transpose_99", "transpose_102"]
        logger.debug("tensor_names: %s", tensor_names)
        tensors = [self.graph.get_tensor_by_name("s3fd/{}:0".format(t_name))
                   for t_name in tensor_names]
        logger.debug("tensors: %s", tensors)
        return tensors

    def set_session(self, target_size, vram_ratio, card_id):
        """ Set the TF Session and initialize """
        # pylint: disable=not-context-manager, no-member
        placeholder = np.zeros((1, 3, target_size[0], target_size[1]))
        config = self.tf.ConfigProto()
        if card_id != -1:
            config.gpu_options.visible_device_list = str(card_id)
        if vram_ratio != 1.0:
            config.gpu_options.per_process_gpu_memory_fraction = vram_ratio

        with self.graph.as_default():
            session = self.tf.Session(config=config)
            self.is_gpu = any("gpu" in str(device).lower() for device in session.list_devices())
            session.run(self.output, feed_dict={self.input: placeholder})
        return session

    def detect_face(self, feed_item):
        """ Detect faces """
        feed_item = feed_item - np.array([104.0, 117.0, 123.0])
        feed_item = feed_item.transpose(2, 0, 1)
        feed_item = feed_item.reshape((1,) + feed_item.shape).astype('float32')
        bboxlist = self.session.run(self.output, feed_dict={self.input: feed_item})
        bboxlist = self.post_process(bboxlist)

        keep = self.nms(bboxlist, 0.3)
        bboxlist = bboxlist[keep, :]
        bboxlist = [x for x in bboxlist if x[-1] >= self.confidence]

        return np.array(bboxlist)

    def post_process(self, bboxlist):
        """ Perform post processing on output """
        retval = list()
        for i in range(len(bboxlist) // 2):
            bboxlist[i * 2] = self.softmax(bboxlist[i * 2], axis=1)
        for i in range(len(bboxlist) // 2):
            ocls, oreg = bboxlist[i * 2], bboxlist[i * 2 + 1]
            stride = 2 ** (i + 2)    # 4,8,16,32,64,128
            poss = zip(*np.where(ocls[:, 1, :, :] > 0.05))
            for _, hindex, windex in poss:
                axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
                score = ocls[0, 1, hindex, windex]
                loc = np.ascontiguousarray(oreg[0, :, hindex, windex]).reshape((1, 4))
                priors = np.array([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]])
                variances = [0.1, 0.2]
                box = self.decode(loc, priors, variances)
                x_1, y_1, x_2, y_2 = box[0] * 1.0
                retval.append([x_1, y_1, x_2, y_2, score])
        retval = np.array(retval)
        if len(retval) == 0:
            retval = np.zeros((1, 5))
        return retval

    @staticmethod
    def softmax(inp, axis):
        """Compute softmax values for each sets of scores in x."""
        return np.exp(inp - logsumexp(inp, axis=axis, keepdims=True))

    @staticmethod
    def decode(loc, priors, variances):
        """Decode locations from predictions using priors to undo
        the encoding we did for offset regression at train time.
        Args:
            loc (tensor): location predictions for loc layers,
                Shape: [num_priors,4]
            priors (tensor): Prior boxes in center-offset form.
                Shape: [num_priors,4].
            variances: (list[float]) Variances of priorboxes
        Return:
            decoded bounding box predictions
        """
        boxes = np.concatenate((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
                                priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])),
                               1)
        boxes[:, :2] -= boxes[:, 2:] / 2
        boxes[:, 2:] += boxes[:, :2]
        return boxes

    @staticmethod
    def nms(dets, thresh):
        """ Perform Non-Maximum Suppression """
        keep = list()
        if len(dets) == 0:
            return keep

        x_1, y_1, x_2, y_2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
        areas = (x_2 - x_1 + 1) * (y_2 - y_1 + 1)
        order = scores.argsort()[::-1]

        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            xx_1, yy_1 = np.maximum(x_1[i], x_1[order[1:]]), np.maximum(y_1[i], y_1[order[1:]])
            xx_2, yy_2 = np.minimum(x_2[i], x_2[order[1:]]), np.minimum(y_2[i], y_2[order[1:]])

            width, height = np.maximum(0.0, xx_2 - xx_1 + 1), np.maximum(0.0, yy_2 - yy_1 + 1)
            ovr = width * height / (areas[i] + areas[order[1:]] - width * height)

            inds = np.where(ovr <= thresh)[0]
            order = order[inds + 1]

        return keep


import numpy as np
from keras.models import load_model
from keras.preprocessing import image
import cv2
import datetime

mymodel = load_model('mymodel.h5')

cap = cv2.VideoCapture(0)
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

while cap.isOpened():
    _, img = cap.read()
    face = face_cascade.detectMultiScale(img, scaleFactor=1.1, minNeighbors=4)
    for (x, y, w, h) in face:
        face_img = img[y:y + h, x:x + w]
        cv2.imwrite('temp.jpg', face_img)
        test_image = image.load_img('temp.jpg', target_size=(150, 150, 3))
        test_image = image.img_to_array(test_image)
        test_image = np.expand_dims(test_image, axis=0)
        pred = mymodel.predict(test_image)[0][0]
        if pred == 1:
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 3)
            cv2.putText(img, 'NO MASK', ((x + w) // 2, y + h + 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 3)
        else:
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 3)
            cv2.putText(img, 'MASK', ((x + w) // 2, y + h + 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
        datet = str(datetime.datetime.now())
        cv2.putText(img, datet, (400, 450), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    cv2.imshow('img', img)

    if cv2.waitKey(1) == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()


import tensorflow as tf
from tensorflow.python.ops import variable_scope as vs
import numpy as np
import pickle
import os,json
from samples import read_clips

from laplace_temporal_net import create_r3d
os.environ['CUDA_VISIBLE_DEVICES'] = "0"

gpus=[0]

video_json_path = 'three_samples.json'
data_root = '/home/pr606/Pictures/ucf_images'
checkpoint_path = '/home/pr606/YUAN/history/tf-R2plus1D/saved_models/model.ckpt'

if __name__ == '__main__':
    tf.logging.set_verbosity(tf.logging.INFO)
    _SAMPLE_VIDEO_FRAMES = 32
    _IMAGE_SIZE = 224
    NUM_CLASS = 101

    with open(video_json_path) as f1:
        annotation_json = json.load(f1)
    data_json = annotation_json["database"]
    label_list = annotation_json["labels"]

    video_name = 'v_GolfSwing_g04_c03'
    video_length = data_json['v_GolfSwing_g04_c03']['annotations']['n_frames']
    labe = data_json['v_GolfSwing_g04_c03']['annotations']['label']
    batchsize_gpu_0 = video_length // _SAMPLE_VIDEO_FRAMES
    frame_indexs = range(video_length)
    video_path = os.path.join(data_root, labe, video_name)

    assignments = [batchsize_gpu_0]
    batch_size = 0
    for n in assignments:
        batch_size += n
    assert batch_size >= 1
    pathes = [video_path]*batchsize_gpu_0

    one_batch = read_clips.get_batch(pathes,[frame_indexs[i*_SAMPLE_VIDEO_FRAMES:(i+1)*_SAMPLE_VIDEO_FRAMES] for i in range(batchsize_gpu_0)])

    data = tf.placeholder(shape=(batch_size, _SAMPLE_VIDEO_FRAMES, _IMAGE_SIZE, _IMAGE_SIZE, 3), dtype=tf.float32,name='clips')
    label = tf.placeholder(shape=(batch_size, NUM_CLASS), dtype=tf.int32,name='labels')

    data_list = tf.split(data,assignments,axis=0,name='sep_0_for_%d_gpus'%len(gpus))
    label_list = tf.split(label,assignments,axis=0,name='sep_1_for_%d_gpus'%len(gpus))

    variable_map = {}
    regular_map = {}
    extra_map = {}
    results = []
    for i in gpus:
        with tf.device("/gpu:%d" % i):
            with tf.variable_scope('parall_model',reuse=tf.AUTO_REUSE):
                result = create_r3d(
                    data=data_list[i],
                    model_depth=18,
                    num_labels=NUM_CLASS,
                    num_input_channels=1,
                    is_decomposed=False,
                    no_bias=1,
                    is_test=1,
                    spatial_bn_mom=0.9,
                    final_spatial_kernel=14,
                    final_temporal_kernel=4,
                )
                if i == 0:
                    # Note: key of variable_map should be specific enough to match the unique variable
                    for var in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
                        print(var.name, var.shape)
                        if 'L101' in var.name:
                            extra_map[var.name[13:].replace(":0", '')] = var
                            continue
                        elif 'conv1_middle/kernel' in var.name or 'conv1/kernel' in var.name:
                            extra_map[var.name[13:].replace(":0", '')] = var
                            continue
                        elif 'conv1_middle_spatbn_relu' in var.name:
                            extra_map[var.name[13:].replace(":0", '')] = var
                            continue
                        elif 'Lp_conv_' in var.name or 'global_step' in var.name:
                            continue
                        elif 'gaussian_filter' in var.name:
                            continue
                        elif 'kernel:0' in var.name:
                            regular_map[var.name[13:].replace(":0", '')] = var
                            continue
                        else:
                            variable_map[var.name[13:].replace(":0", '')] = var
                        labels = tf.stop_gradient(label_list[i],
                                                  name='disallow_grad_labels%d' % i)
                        top1_prediction = tf.nn.in_top_k(result, tf.argmax(labels, 1), 1)
                        top5_prediction = tf.nn.in_top_k(result, tf.argmax(labels, 1), 5)
                else:
                    labels = tf.stop_gradient(label_list[i], name='disallow_grad_labels%d' % i)
                    top1_prediction = tf.concat([top1_prediction, tf.nn.in_top_k(result, tf.argmax(labels, 1), 1)],
                                                axis=0)
                    top5_prediction = tf.concat([top5_prediction, tf.nn.in_top_k(result, tf.argmax(labels, 1), 5)],
                                                axis=0)
                    pass
                tf.get_variable_scope().reuse_variables()
                results.append(result)
                # vs.get_variable_scope().reuse_variables()
    acc_top1 = tf.reduce_mean(tf.cast(top1_prediction, tf.float32))
    acc_top5 = tf.reduce_mean(tf.cast(top5_prediction, tf.float32))
    rgb_resaver = tf.train.Saver(var_list=dict(variable_map, **regular_map, **extra_map), reshape=True)
    nodes = tf.get_collection('source_node')
    features = []
    for var in nodes:
        # parall_model/final_avg/AvgPool3D:0
        # features.append(var)
        if '_com2_GPA' in var.name:
            com2_gpa = var
            continue
        # parall_model_1/final_avg/AvgPool3D:0
        # features.append(var)
        if '_com4_GPA' in var.name:
            com4_gpa = var
            continue
    # assert len(features) != 0
    final_results = results[0] if len(results) == 1 else tf.concat(results, axis=0, name='merge_result_across_gpus')
    # final_features = features[0] if len(features) == 1 else tf.concat(features, axis=0, name='merge_features_across_gpus')

    global_init = tf.global_variables_initializer()
    config = tf.ConfigProto(allow_soft_placement=True)
    # config.gpu_options.allow_growth=True
    config.gpu_options.per_process_gpu_memory_fraction = 0.85
    with tf.Session(config=config) as sess:
        sess.run(global_init)
        feed_dict = {}
        if os.path.exists(checkpoint_path+'.index'):
            rgb_resaver.restore(sess, checkpoint_path)
            print('restoring checkpoint from {} '.format(checkpoint_path))
            tf.logging.info('RGB checkpoint restored')
        else:
            assert False,'can not find any checkpoint files'
        all_activations = {'com2_gpa': None, 'com4_gpa': None, 'mapping': None}
        feed_dict[data] = one_batch  #[:, :, :, :, np.newaxis]
        feat2,feat4 = sess.run([com2_gpa,com4_gpa], feed_dict=feed_dict)

        all_activations['com2_gpa'] = feat2
        all_activations['com4_gpa'] = feat4

        all_activations['mapping'] = video_path

    print(all_activations['com2_gpa'].shape)
    print(all_activations['com4_gpa'].shape)
    with open(video_name+'.pickle', 'wb') as handle:
        pickle.dump(all_activations, handle)


import tensorflow as tf
import numpy as np
#import matplotlib.pyplot as plt
import cv2, os
os.environ["CUDA_VISIBLE_DEVICES"]="1"

from tensorflow.contrib import learn
from Detector.RetinaNet import RetinaNet
from utils.bbox import draw_boxes

mode = learn.ModeKeys.INFER

FLAGS = tf.app.flags.FLAGS

tf.logging.set_verbosity(tf.logging.WARN)
tf.app.flags.DEFINE_string('f', '', 'kernel')
#### Input pipeline
tf.app.flags.DEFINE_integer('input_size', 224,
                            """Input size""")
tf.app.flags.DEFINE_integer('batch_size', 8,
                            """Train batch size""")
tf.app.flags.DEFINE_float('learning_rate', 1e-3,
                            """Learninig rate""")
tf.app.flags.DEFINE_integer('num_input_threads', 4,
                            """Number of readers for input data""")
tf.app.flags.DEFINE_integer('num_classes', 1000,
                            """number of classes""")
tf.app.flags.DEFINE_integer('num_gpus', 2,
                            """The number of gpu""")
tf.app.flags.DEFINE_string('tune_from', 'results_test1/model.ckpt-70000',
                           """Path to pre-trained model checkpoint""")
#tf.app.flags.DEFINE_string('tune_from', 'results_test1/',
#                           """Path to pre-trained model checkpoint""")

#### Train dataset
tf.app.flags.DEFINE_string('train_path', '../data/mjsynth/train',
                           """Base directory for training data""")
tf.app.flags.DEFINE_string('filename_pattern', '*.tfrecord',
                           """File pattern for input data""")
### Validation dataset (during training)
tf.app.flags.DEFINE_string('valid_dataset','VOC',
                          """Validation dataset name""")
tf.app.flags.DEFINE_integer('valid_device', 0,
                           """Device for validation""")
tf.app.flags.DEFINE_integer('valid_batch_size', 8,
                            """Validation batch size""")
tf.app.flags.DEFINE_boolean('use_validation', True,
                            """Whether use validation or not""")
tf.app.flags.DEFINE_integer('valid_steps', 1000,
                            """Validation steps""")

#### Output Path
tf.app.flags.DEFINE_string('output', 'results_test2/model.ckpt-66000',
                           """Directory for event logs and checkpoints""")
#### Training config
tf.app.flags.DEFINE_float('cls_thresh', 0.5,
                            """thresh for class""")
tf.app.flags.DEFINE_float('nms_thresh', 0.5,
                            """thresh for nms""")
tf.app.flags.DEFINE_integer('max_detect', 300,
                            """num of max detect (using in nms)""")
tf.app.flags.DEFINE_string('tune_scope', '',
                           """Variable scope for training""")
tf.app.flags.DEFINE_integer('max_num_steps', 2**21,
                            """Number of optimization steps to run""")
tf.app.flags.DEFINE_boolean('verbose', True,
                            """Print log in tensorboard""")
tf.app.flags.DEFINE_boolean('use_profile', False,
                            """Whether use Tensorflow Profiling""")
tf.app.flags.DEFINE_boolean('use_debug', False,
                            """Whether use TFDBG or not""")
tf.app.flags.DEFINE_integer('save_steps', 2000,
                            """Save steps""")
tf.app.flags.DEFINE_integer('summary_steps', 50,
                            """Save steps""")
tf.app.flags.DEFINE_float('moving_average_decay', 0.9999,
                            """Moving Average dacay factor""")

img_dir = "/root/DB/VOC/VOC2012/JPEGImages/"
VOC = {1 : "motorbike", 2 : "car", 3 : "person", 4 : "bus", 5 : "bird", 6 : "horse", 7 : "bicycle", 8 : "chair", 9 : "aeroplane", 10 : "diningtable", 11 : "pottedplant", 12 : "cat", 13 : "dog", 14 : "boat", 15 : "sheep", 16 : "sofa", 17 : "cow", 18 : "bottle", 19 : "tvmonitor", 20 : "train"}

def get_session_config():
    """Setup session config to soften device placement"""
    config=tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=False)

    return config

def load_pytorch_weight():
    from torch import load

    pt_load = load("weights/resnet50.pth")
    reordered_weights = {}
    pre_train_ops = []

    for key, value in pt_load.items():
        try:
            reordered_weights[key] = value.data.numpy()
        except:
            reordered_weights[key] = value.numpy()

    weight_names = list(reordered_weights)

    #tf_variables = tf.trainable_variables()
    #tf_variables = [v for v in tf_variables if "train_tower_0" and "resnet" in v.name]
    #bn_variables = [v for v in (tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)) if "train_tower_0" and "moving_" in v.name]

    tf_variables = [v for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="train_tower_0") if "resnet" in v.name]
    bn_variables = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="train_tower_0") if "moving_" in v.name]
    fc_variables = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="train_tower_0") if "dense" in v.name]

    print(fc_variables)
    tf_counter = 0
    tf_bn_counter = 0
    tf_fc_counter = 0

    for name in weight_names:
        if "fc" in name:
            pt_assign = reordered_weights[name]
            tf_assign = fc_variables[tf_fc_counter]
            pre_train_ops.append(tf_assign.assign(np.transpose(pt_assign)))
            tf_fc_counter += 1

        elif len(reordered_weights[name].shape) == 4:
            weight_var = reordered_weights[name]
            tf_weight = tf_variables[tf_counter]

            pre_train_ops.append(tf_weight.assign(np.transpose(weight_var, (2, 3, 1, 0))))
            tf_counter += 1

        elif "running_" in name:
            pt_assign = reordered_weights[name]
            tf_assign = bn_variables[tf_bn_counter]

            pre_train_ops.append(tf_assign.assign(pt_assign))
            tf_bn_counter += 1

        else:
            pt_assign = reordered_weights[name]
            tf_assign = tf_variables[tf_counter]

            pre_train_ops.append(tf_assign.assign(pt_assign))
            tf_counter += 1

    return tf.group(*pre_train_ops, name='load_resnet_pretrain')

def _get_init_pretrained(sess):
    saver_reader = tf.train.Saver(tf.global_variables())
    saver_reader.restore(sess, FLAGS.tune_from)

with tf.Graph().as_default():
    image = tf.placeholder(tf.float32, shape=[1, 224, 224, 3], name='image')

    with tf.variable_scope('train_tower_0') as scope:
        net = RetinaNet("resnet50")

        resnet = net.resnet(image, mode)
        gap = tf.nn.avg_pool(resnet[5], ksize=[1, 7, 7, 1], strides=[1, 7, 7, 1], padding='SAME')

        gap = tf.squeeze(gap, [1, 2])
        out = tf.layers.dense(gap, 1000, activation=None)
        out = tf.argmax(out, 1)

    #restore_model = get_init_trained()
    pretrain_op = load_pytorch_weight()
    init_op = tf.group( tf.global_variables_initializer(),
                        tf.local_variables_initializer())

    with tf.Session() as sess:
        sess.run(init_op)
        sess.run(pretrain_op)

        for n, _img in enumerate(os.listdir(img_dir)):
            img = cv2.imread(img_dir + _img)
            img = cv2.resize(img, (224, 224))

            batch_image = np.expand_dims(img, 0)
            print(batch_image.shape)
            result = sess.run(out, feed_dict={image : batch_image})

            #img = draw_boxes(img, box, label)
            #plt.imshow(img)
            print(result)
            if n==1:
                break


import os
import numpy as np
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import rc
rc('font', **{'family': 'serif'})
from matplotlib import rcParams
rcParams['xtick.direction'] = 'out'
rcParams['ytick.direction'] = 'out'
rcParams['xtick.labelsize'] = 18
rcParams['ytick.labelsize'] = 18
rcParams['lines.linewidth'] = 1.85
rcParams['axes.labelsize'] = 20
rcParams.update({'figure.autolayout': True})
np.set_printoptions(threshold=10000)
import pickle

class Grouping(object):
    def __init__(self, G, geoOption, contig, name, bounds, minCutoff, ratioCutoff, groupCutoff, sig_t, E=None):
        self.name = name
        self.G = G
        self.geoOption = geoOption
        self.contig = contig
        self.minCutoff = minCutoff
        self.ratioCutoff = ratioCutoff
        self.groupCutoff = groupCutoff
        self.structure = bounds
        self.mask = np.array([i for c in range(max(bounds) + 1) for i, g in enumerate(bounds) if g == c])
        self.counts = dict(zip(*np.unique(bounds, return_counts=True)))
        self.maxOrder = np.max(list(self.counts.values()))
        self.fname = '{}_{}_{}g_m{}_r{}_g{}'.format(name, geoOption, G, round(minCutoff, 2), round(ratioCutoff, 2), groupCutoff)
        self.sig_t = sig_t
        self.E = E

    def __str__(self):
        s = 'Name: {}'.format(self.name)
        s += '\nGeometry: {}'.format(self.geoOption)
        s += '\nMaxOrder: {}'.format(self.maxOrder)
        #s += '\nStructure: {}'.format(' '.join(self.structure.astype(str)))
        struct = np.concatenate(([0], np.cumsum([self.counts[i] for i in range(max(self.structure) + 1)])))
        s += '\nStructure:'
        for i in range(len(struct) - 1):
            s += '\n\tGroup {}: '.format(i)
            s += '{}-{}'.format(struct[i] + 1, struct[i+1]) if struct[i] + 1 != struct[i+1] else '{}'.format(struct[i] + 1)
        #s += '\nMask: {}'.format(' '.join(self.mask.astype(str)))
        s += '\nCounts: {}'.format(', '.join(['CG-{}={}'.format(key, item) for key, item in self.counts.items()]))
        s += '\nFileName: {}'.format(self.fname)

        return s

    def writeFile(self):
        # Get the coarse group bounds given the total cross sections
        s = '{}\n{}'.format(self.name, self.structure.tolist())

        with open('XS/structure_{}'.format(self.fname), 'w') as f:
            f.write(s)

def getMats(geoOption):
    if geoOption == 'inf':
        return [0]
    elif geoOption == 'full':
        return [0, 3, 9]
    elif geoOption == 'core1':
        return [0, 1, 2, 9]
    elif geoOption == 'core2':
        return [0, 1, 3, 9]
    elif 'c5g7' in geoOption:
        return [0, 3, 9]
        if 'uo2' in geoOption:
            return [5, 9]
        elif 'mox' in geoOption:
            if 'low' in geoOption:
                return [6, 9]
            elif 'mid' in geoOption:
                return [7, 9]
            elif 'high' in geoOption:
                return [8, 9]
            else:
                return [6, 7, 8, 9]
        else:
            return [5, 6, 7, 8, 9]
    else:
        raise NotImplementedError('{} has not been implented'.format(geoOption))

def getXS(fname):
    '''
    Read the input file and extract the total cross sections

    input
        fname: name of the file to read

    output
        sig_t: numpy array containing the total cross sections for each material
    '''

    # Initialize the total cross section list
    sig_t = []
    with open(fname) as f:
        # Get the number of materials and number of groups
        nMat, nGroups = [int(d) for d in f.readline().split()[:2]]
        # Skip 2 lines to pass the energy bounds and velocity
        E = np.array(f.readline().split()).astype(float)[::-1]
        f.readline()

        # Loop over the materials
        for m in range(nMat):
            # Initialize temp list for total XS for material m
            t = []
            # Skip the material name
            f.readline()
            # Get the number of legendre moments
            nLegendre = int(f.readline().split()[0])
            # Loop over groups and extract the total XS
            for g in range(nGroups):
                t.append(float(f.readline().split()[0]))
            # Skip the scattering moments
            for g in range(nGroups * nLegendre):
                f.readline()
            # Append the temp XS list to the main list
            sig_t.append(t)

    return np.array(sig_t), E

def findContiguousBounds(sig_t, minCutoff, ratioCutoff, groupCutoff):
    '''
    Find the coarse group structure given a set of cross sections

    Inputs:
        sig_t: array of the total XS indexed by material [mat, G]
        minCutoff: XS below which the ratio cutoff is ignored
        ratioCutoff: Ratio of largest to smallest XS in a coarse group
        groupCutoff: Maximum number of fine groups in a coarse group
    Outputs:
        bounds: Starting fine group indices for each coarse group
    '''

    def reset(xs):
        '''Reset the min and max cutoffs to the input cross section'''
        minXS = np.maximum(xs, minCutoff)
        maxXS = xs
        return minXS, maxXS

    sig_t = np.array(sig_t)

    # Get the number of groups
    nGroup = len(sig_t[0])

    # Initialize the boundary at the lowest energy group
    bounds = [nGroup + 1]

    # Initialize the cutoff bounds
    minXS, maxXS = reset(sig_t[:,-1])

    # Loop through the cross sections from greatest to least
    for i, xs in enumerate(sig_t.T[::-1]):
        group = nGroup - i

        # Check if the xs is below the min or above the max
        minXS = np.minimum(xs, minXS)
        maxXS = np.maximum(xs, maxXS)
        ratio = maxXS - minXS

        # Check for a ratio that is too large
        if (max(ratio) > ratioCutoff and max(maxXS) > minCutoff) or bounds[-1] - group > groupCutoff:
            bounds.append(group + 1)
            # Reset the cutoff bounds
            minXS, maxXS = reset(xs)

    # Add the highest energy group bound
    bounds.append(1)
    # Reverse to the natural order of structures
    bounds = np.array(bounds[::-1])

    return np.array([i for i, b in enumerate(bounds[1:] - bounds[:-1]) for j in range(b)])

def findNoncontiguousBounds(sig_t, minCutoff=1.0, ratioCutoff=2.0, groupCutoff=60):
    '''
    Find the coarse group structure given a set of cross sections

    Inputs:
        sig_t: array of the total XS indexed by material [mat, G]
        minCutoff: XS below which the ratio cutoff is ignored
        ratioCutoff: Ratio of largest to smallest XS in a coarse group
        groupCutoff: Maximum number of fine groups in a coarse group
    Outputs:
        structure: mapping of fine group to coarse group
    '''

    G = len(sig_t[0])

    # Get the minimum and maximum of all materials
    minXS = np.min(sig_t, axis=0)
    minmap = np.argsort(minXS)[::-1]
    maxXS = np.max(sig_t, axis=0)
    maxmap = np.argsort(maxXS)[::-1]

    unsorted_groups = list(range(G))

    # Initialize a dictionary
    structure = np.zeros(G) - 1
    coarse_group_index = 0
    while len(unsorted_groups) > 0:
        # Get the largest unsorted cross section
        maximum = maxXS[maxmap[0]]
        # If above the minimum cross section cutoff
        if maximum > minCutoff:
            # Find the ratio for the current group
            #cutoff = max(ratioCutoff, maximum / minXS[maxmap[0]])
            cutoff = ratioCutoff
            # Select all groups with a smaller ratio than the cutoff
            cutmap = minmap[maximum - minXS[minmap] <= cutoff]
            if len(cutmap) == 0:
                cutmap = [maxmap[0]]
        else:
            # Select all remaining groups
            cutmap = minmap
        # Only allow groupCutoff groups into the current group
        cutmap = cutmap[:groupCutoff]
        # Assign selected groups to the coarse_group_index
        structure[cutmap] = coarse_group_index
        # Remove the assigned groups from minmap, maxmap, and unsorted_groups
        minmap = np.array([g for g in minmap if g not in cutmap])
        maxmap = np.array([g for g in maxmap if g not in cutmap])
        unsorted_groups = [g for g in unsorted_groups if g not in cutmap]
        # Increment the coarse_group_index
        coarse_group_index += 1

    return (structure * -1 + max(structure)).astype(int)

def computeBounds(G, geoOption, contig, minCutoff, ratioCutoff, groupCutoff):
    '''
    This function determines which fine groups should belong to each coarse group

    The boundaries will not necessarily be contiguous

    Inputs:
        fname: name of the cross section file in anlxs format
    '''

    # Contig options
    names = {0: 'contiguous_min',
             1: 'contiguous_max',
             2: 'contiguous_mean',
             3: 'continuous_difference',
             4: 'noncontiguous_min',
             5: 'noncontiguous_max',
             6: 'noncontiguous_mean',
             7: 'noncontiguous_difference'}

    # Select the material types for the geoOption
    mats = getMats(geoOption)

    # Load the total cross section for group G
    sig_t, E = getXS('XS/{}gXS.anlxs'.format(G))

    # Slice out only the cross sections used for the current geo
    mask = np.array([r in mats for r in range(len(sig_t))]).astype(bool)
    sig_t = sig_t[mask]

    # Number the energy groups
    groups = np.arange(G)

    # Get the minimum XS across all materials
    minXS = np.min(sig_t, axis=0)
    # Get the maximum XS across all materials
    maxXS = np.max(sig_t, axis=0)
    # Get the mean XS across all materials
    meanXS = np.mean(sig_t, axis=0)

    # Select which XS upon which to compute the bounds
    if contig == 0:  # contiguous min
        mask = findContiguousBounds([minXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 1:  # contiguous max
        mask = findContiguousBounds([maxXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 2:  # contiguous mean
        mask = findContiguousBounds([meanXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 3:  # continuous difference
        mask = findContiguousBounds(sig_t, minCutoff, ratioCutoff, groupCutoff)
    elif contig == 4:  # non-contiguous min
        mask = findNoncontiguousBounds([minXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 5:  # non-contiguous max
        mask = findNoncontiguousBounds([maxXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 6:  # non-contiguous mean
        mask = findNoncontiguousBounds([meanXS], minCutoff, ratioCutoff, groupCutoff)
    elif contig == 7:  # non-contiguous difference
        mask = findNoncontiguousBounds(sig_t, minCutoff, ratioCutoff, groupCutoff)

    structure = Grouping(G, geoOption, contig, names[contig], mask, minCutoff, ratioCutoff, groupCutoff, sig_t, E)

    makePlot(structure, minXS, maxXS)

    return structure

def makePlot(structure, minXS, maxXS):
    # Plot the cross sections
    colors = ['b', 'g', 'r', 'c', 'm', 'orange', 'purple', 'grey']
    ls = '-' if structure.contig < 4 else '--'

    groups = np.arange(structure.G)
    mask = structure.structure

    for i, m in enumerate(mask):
        submask = mask == m
        plt.fill_between(groups[submask] + 1, minXS[submask], maxXS[submask], label='Group {}'.format(m), linestyle=ls, color=colors[m%len(colors)])
        #plt.semilogy(groups[i], minXS[i], label='Group {}'.format(m), ls=ls, marker='o', c=colors[m%len(colors)])
        #plt.semilogy(groups[i], maxXS[i], label='Group {}'.format(m), ls=ls, marker='s', c=colors[m%len(colors)])
    plt.yscale('log')
    plt.xlabel('Fine group number')
    plt.ylabel('Total cross section [cm$^{-1}$]')
    plt.xlim([0, structure.G+1])
    plt.grid(True)
    #plt.title(structure.name + ' {}'.format(len(structure.counts.keys())))
    plt.savefig('XS/structure_plot_{}.png'.format(structure.fname), transparent=True)
    plt.clf()

    E = structure.E[::-1]
    E = 0.5 * (E[1:] + E[:-1])

    for i, m in enumerate(mask):
        submask = mask == m
        plt.fill_between(E[submask], minXS[submask], maxXS[submask], label='Group {}'.format(m), linestyle=ls, color=colors[m%len(colors)])
        #plt.semilogy(groups[i], minXS[i], label='Group {}'.format(m), ls=ls, marker='o', c=colors[m%len(colors)])
        #plt.semilogy(groups[i], maxXS[i], label='Group {}'.format(m), ls=ls, marker='s', c=colors[m%len(colors)])
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('Energy [MeV]')
    plt.ylabel('Total cross section [cm$^{-1}$]')
    plt.grid(True)
    #plt.title(structure.name + ' {}'.format(len(structure.counts.keys())))
    plt.savefig('XS/structure_plot_{}_E.png'.format(structure.fname), transparent=True)
    plt.clf()

def getInfo(task, kill=False):
    Gs = [44, 238, 1968]
    geos = ['inf', 'full', 'core1', 'core2', 'c5g7-small']
    contigs = [1]
    mins = [0]
    ratios = [1.3]
    groups = [60]

    item = 0
    groupBounds = []

    for geo in geos:
        for G in Gs:
            for contig in contigs:
                for min_cutoff in mins:
                    for ratio_cutoff in ratios:
                        for group_cutoff in groups:
                            if item == task and not kill:
                                return G, geo, contig, min_cutoff, ratio_cutoff, group_cutoff
                            else:
                                item += 1
        groupBounds.append(item)
    if kill:
        groupBounds = [0] + groupBounds
        for i, G in enumerate(Gs):
            print('Group {:4d} cutoffs: {:>5d}-{:<5d}'.format(G, groupBounds[i] + 1, groupBounds[i+1]))
    exit()

def f(s):
    '''Convert a structure into a proper string'''
    try:
        return '{}'.format(s[0] + 1) if s[1] - 1 == s[0] else '{}-{}'.format(s[0] + 1, s[1])
    except IndexError:
        return ''

if __name__ == '__main__':
    # getInfo(1, True)

    #task = os.environ['SLURM_ARRAY_TASK_ID']
    #task = int(task) - 1

    for geo in ['full', 'core1', 'core2']:
        print(geo)
        s = '& 44-group & \multicolumn{2}{c|}{238-group} & \multicolumn{4}{c}{1968-group}\\\\\n'
        s += '\hline\n'
        s += 'add\'l. \# & (+0) & (+0) & (+20) & (+0) & (+20) & (+40) & (+60)\\\\\n'
        s += '\hline\n'

        # Build the coarse-group structure
        s44 = computeBounds(44, geo, 1, 0.0, 1.3, 60)
        c44 = np.concatenate(([0], np.cumsum([s44.counts[i] for i in range(max(s44.structure)+1)])))
        s238 = computeBounds(238, geo, 1, 0.0, 1.3, 60)
        c238 = np.concatenate(([0], np.cumsum([s238.counts[i] for i in range(max(s238.structure)+1)])))
        s1968 = computeBounds(1968, geo, 1, 0.0, 1.3, 60)
        c1968 = np.concatenate(([0], np.cumsum([s1968.counts[i] for i in range(max(s1968.structure)+1)])))

        for i in range(20):
            s += 'CG {:2d} & {:5s} & {:7s} & {:7s} & {:9s} & {:9s} & {:9s} & {:9s} \\\\\n'.format(i+1, f(c44[i:i+2]), f(c238[i:i+2]), f(c238[i+20:i+22]), f(c1968[i+0:i+2]), f(c1968[i+20:i+22]), f(c1968[i+40:i+42]), f(c1968[i+60:i+62]))
        print(s)

        #pickle.dump(structure, open('XS/structure{}_{}_{}g_m{}_r{}_g{}.p'.format(contig, geo, G, min_cutoff, ratio_cutoff, group_cutoff), 'wb'))


import os
import cv2
import sys
import time
import math
import json
import glob
import argparse
import matplotlib
import numpy as np
import tensorflow as tf
import tensorlayer as tl
from hyperpose import Config,Model,Dataset

os.environ['CUDA_VISIBLE_DEVICES']='0,1'
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[0], True)
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[1], True)

model_type = "Openpose"
# model_backbone = "Default"
# model_backbone = "Mobilenetv2"
model_backbone = "Resnet18"
# MobilenetV1.MobilenetV2.Vggtiny.Vgg16.Vgg19.Resnet18.Resnet50

dataset_type = "USERDEF"
model_name = dataset_type+'_'+model_type+'_' + model_backbone + '_'+time.strftime("%Y%m%d_%H%M")
dataset_version ="2020"
dataset_path = "./data"
train_type = "Single_train"
kf_optimizer ='Sync_avg'
userdef_dataset = 1
useradd_data_path = "./data"

domainadapt_data_path = None

optim_type ="Adam"
learning_rate =1e-4
log_interval =1e2
save_interval =5e3

Config.set_model_name(model_name)
Config.set_model_type(Config.MODEL[model_type])
Config.set_model_backbone(Config.BACKBONE[model_backbone])
Config.set_log_interval(log_interval)
Config.set_save_interval(save_interval)
# config train
Config.set_train_type(Config.TRAIN[train_type])
Config.set_learning_rate(learning_rate)
Config.set_optim_type(Config.OPTIM[optim_type])
Config.set_kungfu_option(Config.KUNGFU[kf_optimizer])

# config dataset
print(f"test enabling dataset:{userdef_dataset}")
Config.set_userdef_dataset(userdef_dataset)
Config.set_dataset_type(Config.DATA[dataset_type])
Config.set_dataset_path(dataset_path)
Config.set_dataset_version(dataset_version)

print('\n' * 4)

# train
config = Config.get_config()

config.pretrain.batch_size=4
config.eval.batch_size=1
config.train.batch_size=4
config.model.hin = 48*2
config.model.win = 64*2
config.model.hout = 6*2 #48/8
config.model.wout = 8*2 #64/8
config.train.save_interval = 5000
config.train.n_step = 60000
# config.train.lr_decay_every_step = 2000

print(config.model)
print('\n' * 2)
print(config.train)
print('\n' * 2)
print(config.data)
print('\n' * 4)

model = Model.get_model(config)

train = Model.get_train(config)

dataset = Dataset.get_dataset(config)

train(model, dataset)


import torch
import torch.nn.functional as F
from torch.distributions import Categorical
import numpy as np
import os

class Agent:
    def __init__(self, p_net, v_net, optim_p, optim_v, device):
        super(Agent, self).__init__()
        self.p_net = p_net
        self.v_net = v_net
        self.optim_p = optim_p
        self.optim_v = optim_v
        self.device = device

    def get_action(self, obs):
        obs = obs.long()
        with torch.no_grad():
            pi = self.p_net(obs)
            value = self.v_net(obs)
        c = Categorical(pi)
        action = c.sample()
        # Action, Value, llh(対数尤度(log likelihood))
        return action.squeeze().item(), value.item(), c.log_prob(action).item(), pi.detach().cpu().numpy()

    def get_llhs(self, obs, actions):
        with torch.no_grad():
            pi = self.p_net(obs)
            llhs = Categorical(pi).log_prob(actions)
        return llhs.detach()

    def update(self, traj, batch_size, epochs, clip_param, v_loss_c, ent_c, max_grad_norm):
        iterate = traj.iterate(batch_size, epochs)
        loss_vals = []
        for batch in iterate:
            loss_vals.append(self._update_batch(batch, clip_param, v_loss_c, ent_c, max_grad_norm))

        loss_vals = np.mean(loss_vals, axis=0)
        return loss_vals

    def save_model(self, filename, info, suffix):
        if suffix:
            filename += "_" + suffix
        filename += ".pkl"
        print("save model: {}".format(filename))
        save_data = {"state_dict_p": self.p_net.state_dict(), "state_dict_v": self.v_net.state_dict(), "info": info}
        torch.save(save_data, filename)

    def load_model(self, filename, suffix):
        if suffix:
            filename += "_" + suffix
        filename += ".pkl"
        if not os.path.exists(filename):
            return None
        print("load model: {}".format(filename))
        load_data = torch.load(filename, map_location=self.device)
        self.p_net.load_state_dict(load_data["state_dict_p"])
        self.v_net.load_state_dict(load_data["state_dict_v"])
        return load_data["info"]

    def train(self):
        self.p_net.train()
        self.v_net.train()

    def eval(self):
        self.p_net.eval()
        self.v_net.eval()

    def _update_batch(self, batch: dict, clip_param: float, v_loss_c: float, ent_c: float, max_grad_norm: float):
        obs = batch["obs"].long()
        actions = batch["actions"]
        llhs = batch["llhs"]
        advantages = batch["advantages"]
        returns = batch["returns"]

        pol_loss, entropy = self._update_pol(obs, actions, llhs, advantages, clip_param, ent_c, max_grad_norm)
        v_loss = self._update_val(obs, returns, v_loss_c, max_grad_norm)

        return pol_loss, v_loss, entropy

    def _update_pol(self, obs, actions, llhs, advantages, clip_param, ent_c, max_grad_norm):
        # Policy
        pi = self.p_net(obs)
        c = Categorical(pi)
        new_llhs = c.log_prob(actions)
        ratio = torch.exp(new_llhs - llhs)  # pi(a|s) / pi_old(a|s)
        clip_ratio = ratio.clamp(1.0 - clip_param, 1.0 + clip_param)
        pol_loss = (torch.max(-ratio * advantages, -clip_ratio * advantages)).mean()
        # Entropy
        entropy = -ent_c * c.entropy().mean()

        total_loss = pol_loss + entropy
        self.optim_p.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.p_net.parameters(), max_grad_norm)
        self.optim_p.step()

        return pol_loss.item(), entropy.item()

    def _update_val(self, obs, returns, v_loss_c, max_grad_norm):
        # Value
        value = self.v_net(obs)
        v_loss = v_loss_c * F.smooth_l1_loss(value.squeeze(1), returns).mean()

        self.optim_v.zero_grad()
        v_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.v_net.parameters(), max_grad_norm)
        self.optim_v.step()

        return v_loss.item()

class Discriminator:
    def __init__(self, pseudo_rew_net, shaping_val_net, optim_discrim, device):
        self.pseudo_rew_net = pseudo_rew_net
        self.shaping_val_net = shaping_val_net
        self.optim_discrim = optim_discrim
        self.device = device

    def get_pseudo_reward(self, obs):
        obs = obs.long()
        with torch.no_grad():
            reward = self.pseudo_rew_net(obs)
        return reward.item()

    def update(self, agent_traj, expert_traj, batch_size, gamma, agent):
        agent_iterate = agent_traj.iterate(batch_size, epoch=1)
        expert_iterate = expert_traj.iterate(batch_size, epoch=1)
        loss_vals = []
        for agent_batch, expert_batch in zip(agent_iterate, expert_iterate):
            loss_vals.append(self._update_batch(agent_batch, expert_batch, gamma, agent))

        loss_vals = np.mean(loss_vals, axis=0)
        return loss_vals

    def save_model(self, filename, suffix):
        if suffix:
            filename += "_" + suffix
        filename += ".pkl"
        print("save model: {}".format(filename))
        save_data = {
            "state_dict_rew": self.pseudo_rew_net.state_dict(),
            "state_dict_val": self.shaping_val_net.state_dict(),
        }
        torch.save(save_data, filename)

    def load_model(self, filename, suffix):
        if suffix:
            filename += "_" + suffix
        filename += ".pkl"
        if not os.path.exists(filename):
            return None
        print("load model: {}".format(filename))
        load_data = torch.load(filename, map_location=self.device)
        self.pseudo_rew_net.load_state_dict(load_data["state_dict_rew"])
        self.shaping_val_net.load_state_dict(load_data["state_dict_val"])

    def train(self):
        self.pseudo_rew_net.train()
        self.shaping_val_net.train()

    def eval(self):
        self.pseudo_rew_net.eval()
        self.shaping_val_net.eval()

    def _update_batch(self, agent_batch, expert_batch, gamma, agent):
        agent_loss = self._calc_airl_loss(agent_batch, gamma, agent, True)
        expert_loss = self._calc_airl_loss(expert_batch, gamma, agent, False)
        discrim_loss = agent_loss + expert_loss

        self.optim_discrim.zero_grad()
        discrim_loss.backward()
        self.optim_discrim.step()

        return agent_loss.item(), expert_loss.item()

    def _calc_airl_loss(self, batch, gamma, agent, is_agent):
        obs = batch["obs"].long()
        actions = batch["actions"]
        next_obs = batch["next_obs"].long()
        dones = batch["dones"]

        value = self.shaping_val_net(obs).squeeze(1)
        next_value = self.shaping_val_net(next_obs).squeeze(1)
        reward = self.pseudo_rew_net(obs).squeeze(1)
        energies = reward + (1 - dones) * gamma * next_value - value
        llhs = agent.get_llhs(obs, actions)
        logits = energies - llhs

        if is_agent:
            target = torch.zeros(len(logits)).to(self.device)
        else:
            target = torch.ones(len(logits)).to(self.device)
        discrim_loss = 0.5 * F.binary_cross_entropy_with_logits(logits, target)
        return discrim_loss