import config as cfg
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array

class predict_emotions():
    def __init__(self):
        # load emotion detection model
        self.model = load_model(cfg.path_model)

    def preprocess_img(self,face_image,rgb=True,w=48,h=48):
        face_image = cv2.resize(face_image, (w,h))
        if rgb == False:
            face_image = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY)
        face_image = face_image.astype("float") / 255.0
        face_image= img_to_array(face_image)
        face_image = np.expand_dims(face_image, axis=0)
        return face_image

    def get_emotion(self,img,boxes_face):
        emotions = []
        if len(boxes_face)!=0:
            for box in boxes_face:
                y0,x0,y1,x1 = box
                face_image = img[x0:x1,y0:y1]
                # preprocesar data
                face_image = self.preprocess_img(face_image ,cfg.rgb, cfg.w, cfg.h)
                # predecir imagen
                prediction = self.model.predict(face_image)
                emotion = cfg.labels[prediction.argmax()]
                emotions.append(emotion)
        else:
            emotions = []
            boxes_face = []
        return boxes_face,emotions


import torch
from torch.autograd import Variable
import torch.nn as nn
import numpy as np
import cv2

from gesture.environments.utils import rgb_render, rgb_tensor_render

# === Social ===
def exploration(pi, current, targets, rollouts, args, result,  env):
    ''' Exploration part of PPO training:
    1. Sample actions and gather rewards trajectory for num_steps.
    2. Reset states and rewards if some environments are done.
    3. Keep track of means and std fo
    visualizing progress.
    '''
    stds = []
    s, st, o, ot = current()
    for step in range(args.num_steps):
        # add step count
        pi.n += 1

        # Sample actions
        value, action, action_log_prob, a_std = pi.sample(s, st, o, ot)
        cpu_actions = action.data.squeeze(1).cpu().numpy()

        # Observe reward and next state
        state, s_target, obs, o_target, reward, done, info = env.step(cpu_actions)
        reward = torch.from_numpy(reward).view(args.num_proc, -1).float()
        masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
        result.episode_rewards += reward

        if args.render:
            env.render()

        if sum(done) > 0:
            # Clear episode reward and update final rewards
            result.tmp_final_rewards *= masks
            result.tmp_final_rewards += (1 - masks) * result.episode_rewards
            result.episode_rewards *= masks
            result.update_list()

            env.set_target(targets())

        if args.cuda:
            masks = masks.cuda()

        # Reset current states for envs done and
        # update current state and add data to rollouts
        current.check_and_reset(masks)
        current.update(state, s_target, obs, o_target)
        s, st, o, ot = current()
        rollouts.insert(step, s, st, o, ot,
                        action.data,
                        action_log_prob.data,
                        value.data,
                        reward,
                        masks)

def train(pi, args, rollouts, optimizer_pi):
    last_value, _, _, _ = pi.sample(*rollouts.get_last())
    rollouts.compute_returns(last_value.data, args.no_gae, args.gamma, args.tau)

    # Calculate Advantage (normalize)
    advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

    vloss, ploss, ent = 0, 0, 0
    for e in range(args.ppo_epoch):
        data_generator = rollouts.Batch(advantages, args.batch_size)
        for sample in data_generator:

            states_batch, state_target_batch, obs_batch, \
            obs_target_batch, actions_batch, return_batch, \
                masks_batch, old_action_log_probs_batch, adv_targ = sample

            v, a_logprobs, entro = pi.evaluate_actions(states_batch,
                                                       state_target_batch,
                                                       obs_batch,
                                                       obs_target_batch,
                                                       actions_batch)
            # PPO loss
            adv_targ = Variable(adv_targ)
            ratio = torch.exp(a_logprobs - Variable(old_action_log_probs_batch))
            surr1 = ratio * adv_targ
            surr2 = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) * adv_targ
            action_loss = -torch.min(surr1, surr2).mean()  # PPO's pessimistic surrogate (L^CLIP)
            value_loss = (Variable(return_batch) - v).pow(2).mean()

            # update
            optimizer_pi.zero_grad()
            (value_loss+action_loss- entro *args.entropy_coef).backward()
            nn.utils.clip_grad_norm(pi.parameters(), args.max_grad_norm)
            optimizer_pi.step()

            vloss += value_loss
            ploss += action_loss.abs()
            ent += entro

    vloss /= args.ppo_epoch
    ploss /= args.ppo_epoch
    ent /= args.ppo_epoch
    # return value_loss, action_loss, dist_entropy
    return vloss, ploss, ent

def trainAll(pi, args, rollouts, optimizer_pi, U_loss=None):
    last_value, _, _, _ = pi.sample(*rollouts.get_last())
    rollouts.compute_returns(last_value.data, args.no_gae, args.gamma, args.tau)

    # Calculate Advantage (normalize)
    advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

    vloss, ploss, ent, uloss = 0, 0, 0, 0
    for e in range(args.ppo_epoch):
        data_generator = rollouts.Batch(advantages, args.batch_size)
        for sample in data_generator:
            optimizer_pi.zero_grad()

            states_batch, state_target_batch, obs_batch, \
            obs_target_batch, actions_batch, return_batch, \
                masks_batch, old_action_log_probs_batch, adv_targ = sample

            # Reshape to do in a single forward pass for all steps
            v, a_logprobs, entro, st_pred = pi.evaluate_actions(states_batch,
                                                                state_target_batch,
                                                                obs_batch,
                                                                obs_target_batch,
                                                                actions_batch)
            # PPO loss
            adv_targ = Variable(adv_targ)
            ratio = torch.exp(a_logprobs - Variable(old_action_log_probs_batch))
            surr1 = ratio * adv_targ
            surr2 = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) * adv_targ
            action_loss = -torch.min(surr1, surr2).mean()  # PPO's pessimistic surrogate (L^CLIP)
            value_loss = (Variable(return_batch) - v).pow(2).mean()

            # Understand loss
            understand_loss = U_loss(st_pred, Variable(state_target_batch, requires_grad=False))

            # update
            (understand_loss+value_loss + action_loss - entro * args.entropy_coef).backward()
            nn.utils.clip_grad_norm(pi.parameters(), args.max_grad_norm)
            optimizer_pi.step()

            vloss += value_loss
            uloss += understand_loss
            ploss += action_loss.abs()
            ent += entro

    vloss /= args.ppo_epoch
    uloss /= args.ppo_epoch
    ploss /= args.ppo_epoch
    ent /= args.ppo_epoch
    # return value_loss, action_loss, dist_entropy
    return vloss, ploss, ent, uloss

# Pepper
def explorationPepper(pi, current, targets, rollouts, args, result,  env):
    ''' Exploration part of PPO training:
    1. Sample actions and gather rewards trajectory for num_steps.
    2. Reset states and rewards if some environments are done.
    3. Keep track of means and std fo
    visualizing progress.
    '''
    stds = []
    s, st, _, _ = current()
    for step in range(args.num_steps):
        # add step count
        pi.n += 1

        # Sample actions
        value, action, action_log_prob, a_std = pi.sample(s, st)
        cpu_actions = action.data.squeeze(1).cpu().numpy()

        # Observe reward and next state
        state, o = env.step(cpu_actions)
        reward = torch.from_numpy(reward).view(args.num_proc, -1).float()
        masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
        result.episode_rewards += reward

        if args.render:
            env.render()

        if sum(done) > 0:
            # Clear episode reward and update final rewards
            result.tmp_final_rewards *= masks
            result.tmp_final_rewards += (1 - masks) * result.episode_rewards
            result.episode_rewards *= masks
            result.update_list()

        if args.cuda:
            masks = masks.cuda()

        # Reset current states for envs done and
        # update current state and add data to rollouts
        current.check_and_reset(masks)
        current.update(state, s_target, obs, o_target)
        s, st, o, ot = current()
        rollouts.insert(step, s, st, o, o,
                        action.data,
                        action_log_prob.data,
                        value.data,
                        reward,
                        masks)

def trainPepper(pi, args, rollouts, optimizer_pi):
    last_value, _, _, _ = pi.sample(*rollouts.get_last())
    rollouts.compute_returns(last_value.data, args.no_gae, args.gamma, args.tau)

    # Calculate Advantage (normalize)
    advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

    vloss, ploss, ent = 0, 0, 0
    for e in range(args.ppo_epoch):
        data_generator = rollouts.Batch(advantages, args.batch_size)
        for sample in data_generator:

            states_batch, state_target_batch, obs_batch, \
            obs_target_batch, actions_batch, return_batch, \
                masks_batch, old_action_log_probs_batch, adv_targ = sample

            # Reshape to do in a single forward pass for all steps
            v, a_logprobs, entro = pi.evaluate_actions(states_batch,
                                                       state_target_batch,
                                                       actions_batch)
            # PPO loss
            adv_targ = Variable(adv_targ)
            ratio = torch.exp(a_logprobs - Variable(old_action_log_probs_batch))
            surr1 = ratio * adv_targ
            surr2 = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) * adv_targ
            action_loss = -torch.min(surr1, surr2).mean()  # PPO's pessimistic surrogate (L^CLIP)
            value_loss = (Variable(return_batch) - v).pow(2).mean()

            # update
            optimizer_pi.zero_grad()
            (value_loss+action_loss- entro *args.entropy_coef).backward()
            nn.utils.clip_grad_norm(pi.parameters(), args.max_grad_norm)
            optimizer_pi.step()

            vloss += value_loss
            ploss += action_loss.abs()
            ent += entro

    vloss /= args.ppo_epoch
    ploss /= args.ppo_epoch
    ent /= args.ppo_epoch
    # return value_loss, action_loss, dist_entropy
    return vloss, ploss, ent


from sklearn.metrics import f1_score,recall_score,precision_score,confusion_matrix,accuracy_score

import torch
import torch.nn as nn
import copy
import random
import pandas as pd
import numpy as np
from tqdm import trange
import pickle
import json
import sys
import time
from sklearn.model_selection import train_test_split

sys.path.append("classes")
from loss_functions import NuLogsyLossCompute
from model import *
from networks import *
from tokenizer import *
from data_loader import *
from prototype import get_prototypes
import torch.nn.functional as F
import spacy

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier

def convert_normal_anomaly(x):
    if x == 'trace':
        return 'trace'
    elif x == 'warn':
        return 'warning'
    elif x == 'warning':
        return 'warning'
    elif x == 'info':
        return 'normal'
    elif x == 'debug':
        return 'debug'
    elif x == 'log':
        return 'log'
    else:
        return 'anomaly'

def convert_error_info(x):
    if x == 'trace':
        return 'trace'
    elif x == 'warn':
        return 'log'
    elif x == 'warning':
        return 'log'
    elif x == 'info':
        return 'info'
    elif x == 'debug':
        return 'debug'
    elif x == 'log':
        return 'log'
    else:
        return 'error'

def convert_error_warning(x):
    if x == 'trace':
        return 'trace'
    elif x == 'warn':
        return 'warning'
    elif x == 'warning':
        return 'warning'
    elif x == 'info':
        return 'debug'
    elif x == 'debug':
        return 'debug'
    elif x == 'log':
        return 'log'
    else:
        return 'error'

def convert_info_warning(x):
    if x == 'trace':
        return 'trace'
    elif x == 'warn':
        return 'warning'
    elif x == 'warning':
        return 'warning'
    elif x == 'info':
        return 'info'
    elif x == 'debug':
        return 'debug'
    elif x == 'log':
        return 'log'
    else:
        return 'log'

def convert_error_info_warning(x):
    if x == 'trace':
        return 'trace'
    elif x == 'warn':
        return 'warning'
    elif x == 'warning':
        return 'warning'
    elif x == 'info':
        return 'info'
    elif x == 'debug':
        return 'debug'
    elif x == 'log':
        return 'log'
    else:
        return 'error'

def read_data(path):
    print("Reading data at path ", path)
    return pd.read_csv(path).drop(columns=["Unnamed: 0"])

def preprocess_data(df, scenario, verbose=True):

    if verbose:
        print("Filtering the special characters in the dataframe!")
    df['log_message'] = df['log_message'].str.replace("\<\*\>", " ")
    df['log_message'] = df['log_message'].str.replace("\[STR\]", " ")
    df['log_message'] = df['log_message'].str.replace("\[NUM\]", " ")
    if verbose:
        print("Converting the classes into required categories. Pair or triplet of (INFO, ERROR, WARNING). ")

    if scenario=="error_warning":
        df.loc[:, 'log_level'] = df.loc[:, 'log_level'].apply(lambda x: convert_error_warning(x))
    elif scenario == "info_warning":
        df.loc[:, 'log_level'] = df.loc[:, 'log_level'].apply(lambda x: convert_info_warning(x))
    elif scenario == "info_error":
        df.loc[:, 'log_level'] = df.loc[:, 'log_level'].apply(lambda x: convert_error_info(x))
    elif scenario=="info_error_warning":
        df.loc[:, 'log_level'] = df.loc[:, 'log_level'].apply(lambda x: convert_error_info_warning(x))
    else:
        print("Insert a valid scenario, one in error_warning, info_warning, info_error")
        exit(-1)

    if verbose:
        print("Prior removing (DEBUG, LOG and TRACE) ", df.shape)
    df = df[df['log_level'] != 'debug']
    df = df[df['log_level'] != 'log']
    df = df[df['log_level'] != 'trace']

    if verbose:
        print("Size after removal ", df.shape)
    indecies_to_preserve = df.index
    df = df.reset_index()
    df = df.drop("index", axis=1)
    return df, indecies_to_preserve

def extract_load(df):
    print("Split descriptive and target data into numpay arrays.")
    load = df['log_messages'].values
    labels = df['label'].values
    return load, labels

def tokenization_dataset(df, load, labels, label_mapper):
    tokenizer = LogTokenizer()
    tokenized = []
    for i in trange(0, len(df)):
        tokenized.append(np.array(tokenizer.tokenize(df['log_messages'][i])))
    labels_tokenized = [label_mapper[label] for label in labels]
    return tokenized, labels_tokenized, tokenizer

def run_train(dataloader, model, optimizer, f_loss, epoch, polars=None, device="cpu"):
    model.train()
    total_loss = 0
    start = time.time()
    for i, batch in enumerate(dataloader):
        load, y = batch
        if polars is not None:
            y = polars[y.numpy()]
            y = torch.autograd.Variable(y).cuda()
        if device=="gpu":
            out = model.forward(load.cuda().long(), None)
        else:
            out = model.forward(load.long(), None)
        if isinstance(f_loss, nn.CosineSimilarity):
            loss = (1 - f_loss(out, y)).pow(2).sum()
        else:
            if device == "gpu":
                loss = f_loss(out, y.cuda().long())
            else:
                loss = f_loss(out, y.long())
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        total_loss += loss
        elapsed = time.time() - start
        if i % 5 == 0:
            print("Epoch %d Train Step: %d / %d Loss: %f" %
                  (epoch, i, len(dataloader), loss), end='\r')

    print("Epoch %d Train Step: %d / %d Loss: %f" %
          (epoch, i, len(dataloader), loss), end='\r')
    return total_loss / len(dataloader)

def run_test(dataloader, model, optimizer, f_loss, epoch, polars=None, device="cpu"):
    model.eval()
    preds = []
    # tmps = []
    scores_head1 = []
    scores_head2 = []
    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            load, y = batch
            if device=="gpu":
                out = model.forward(load.cuda().long(), None)
            else:
                out = model.forward(load.long(), None)

            if device=="cuda":
                tmp = out.detach().cpu().numpy()
            else:
                tmp = out.detach().cpu().numpy()
            preds += list(np.argmax(tmp, axis=1))

            # tmps += list(tmp)
            # scores_head1 += model.encoder.layers[0].self_attn.attn[:, 0, :, :].detach().cpu()
            # scores_head2 += model.encoder.layers[0].self_attn.attn[:, 1, :, :].detach().cpu()
    return preds, scores_head1, scores_head2

def run_optimizer(model, train_dataloader, test_dataloader, labels_test, optimizer, n_epochs, f_loss, polars, class_weights, device):
    conf_matrix_good = []
    best_f1 = 0
    best_preds = 0
    epoch_idx = 0
    for epoch in range(1, 1 + n_epochs):

        loss = run_train(train_dataloader, model, optimizer, f_loss, epoch, polars, device=device)

        preds1, scores11, scores12 = run_test(test_dataloader, model, optimizer, f_loss, epoch, polars, device=device)
        # if epoch % 5 == 0:
        print("Epoch", epoch)
        print("Epoch %d Train Loss: %f" % (epoch, loss), " " * 30)
        print("----------GOOD REPOS----------")
        print(f"Accuracy:{round(accuracy_score(preds1, labels_test), 2)}")
        print(f"f1_score:{round(f1_score(preds1, labels_test, average='binary'), 2)}")
        print(f"recall_score:{round(recall_score(preds1, labels_test, average='binary'), 2)}")
        print(f"precision_score:{round(precision_score(preds1, labels_test, average='binary'), 2)}")
        print(confusion_matrix(preds1, labels_test))
        f1sc = f1_score(preds1, labels_test, average='binary')

        if f1sc > best_f1:
            best_f1 = f1sc
            epoch_idx = epoch-1
            best_preds = preds1

        conf_matrix_good.append(confusion_matrix(preds1, labels_test))
    return model, best_preds, conf_matrix_good, best_f1, epoch_idx

def merge_data(list_paths):
    tmp_list = []
    list_indecies_test = []
    for idx, fname in enumerate(list_paths):
        tmp_df = pd.read_csv(fname)
        tmp_df["version"] = np.ones(tmp_df.shape[0])*(idx+1)
        if idx==0:
            list_indecies_test.append(tmp_df.shape[0])
        else:
            list_indecies_test.append(list_indecies_test[-1] + tmp_df.shape[0])
        tmp_list.append(tmp_df)
    return pd.concat(tmp_list, axis=0), list_indecies_test

def split_train_test_dfs(df, train_versions, test_versions):
    train_dfs = []
    for version in train_versions:
        train_dfs.append(df[df.version==version].index)
    test_dfs = []
    for version in test_versions:
        test_dfs.append(df[df.version==version].index)
    return np.hstack(train_dfs), np.hstack(test_dfs)

def select_relevant_logs(repo, top_ranked_repos):
    if repo in top_ranked_repos:
        return True
    else:
        return False

def word2_vec_representation(df, load, labels, label_mapper, nlp):
    tokenizer = LogTokenizer()
    tokenized = []
    for i in trange(0, len(df)):
        tokenized.append(nlp(df['log_messages'][i]).vector)
    labels_tokenized = [label_mapper[label] for label in labels]
    return tokenized, labels_tokenized, tokenizer

from collections import defaultdict

def evaluate(preds1, load_test_good_repos_labels):
    print("********"*10)
    print("----------GOOD REPOS----------")

    print(f"Accuracy:{round(accuracy_score(preds1, load_test_good_repos_labels), 2)}")
    print(f"f1_score:{round(f1_score(preds1, load_test_good_repos_labels, average='binary'), 2)}")
    print(f"recall_score:{round(recall_score(preds1, load_test_good_repos_labels, average='binary'), 2)}")
    print(f"precision_score:{round(precision_score(preds1, load_test_good_repos_labels, average='binary'), 2)}")
    print(confusion_matrix(preds1, load_test_good_repos_labels))

    d = {}
    d["Accuracy"] = accuracy_score(preds1, load_test_good_repos_labels)
    d['F1_score'] = f1_score(preds1, load_test_good_repos_labels, average='binary')
    d["recall_score"] = recall_score(preds1, load_test_good_repos_labels, average='binary')
    d["precision_score"] = precision_score(preds1, load_test_good_repos_labels, average='binary')
    d["confusion_matrix"] = confusion_matrix(preds1, load_test_good_repos_labels)

    return d

nlp = spacy.load("en_core_web_sm")

def extract_rules_train(data):
    stats_ = defaultdict(dict)
    log_messages = data.log_messages
    for idx, doc in enumerate(log_messages):
        doc = nlp(doc)
        stri = []
        for token in doc:
            stri.append(token.dep_)
        stats_[idx] = "_".join(sorted(dict.fromkeys(stri)))
        if idx % 1000 == 0:
            print(idx)
    q = pd.DataFrame(stats_, index=[0]).T
    q = q.reset_index()

    unique_rules = np.unique(q.iloc[:, 1])
    return unique_rules

def predict(data, train_rules):
    stats_ = defaultdict(dict)
    log_messages = data.log_messages
    for idx, doc in enumerate(log_messages):
        doc = nlp(doc)
        stri = []
        for token in doc:
            stri.append(token.dep_)
        stats_[idx] = "_".join(sorted(dict.fromkeys(stri)))
        if idx % 1000 == 0:
            print(idx)
    q = pd.DataFrame(stats_, index=[0]).T
    q = q.reset_index()
    predictions = []

    for idx in range(q.shape[0]):
        flag = True
        for rule in train_rules:
            if rule == q.iloc[idx, 1]:
                predictions.append(0)
                flag = False
                break
        if flag == True:
            predictions.append(1)
    return predictions

def syntactic_rule(train_data, test_data, test_labels):
    train_data = train_data[train_data.label==0]
    train_rules = extract_rules_train(train_data)
    test_prediction = predict(test_data, train_rules)
    test_data["syntactic_prediction"] = test_prediction
    return test_labels, np.array(test_prediction), test_data

####### USE TO GENERATE TOP RANKED REPOS
# PATH_v1= "../../3_preprocessed_data/learning_scenario_3/quality_training_dataset_3.csv"
# PATH_filterd = "../../3_preprocessed_data/filtered_log_df.csv"
# PATH_ranks = "../../3_preprocessed_data/stars_repos.csv"
#
# df = pd.read_csv(PATH_v1)
# ranks = pd.read_csv(PATH_ranks)
# top_ranked_repos = ranks.repo_name.iloc[:100].values
#
# filtered_logs = pd.read_csv(PATH_filterd).repo_link.loc[df.key]
# filtered_logs = filtered_logs.reset_index()
# filtered_logs.columns = ["key", "repo_link"]
# df = pd.merge(df, filtered_logs, on="key", how="inner")
#
# df['to_drop'] = df.repo_link.apply(lambda x: select_relevant_logs(x, top_ranked_repos))
# df = df[df.to_drop==True]
#
#
#
#
#

# scenario = "info_error_warning" # ONE IN: "info_warning", "info_error", "error_warning", "info_error_warning"

def preprocess_log_msg(x):
    x = x.replace(">", "")
    x = x.replace("<", "")
    x = x.replace("*", "")
    x = x.replace("[", "")
    x = x.replace("]", "")
    return x

###################  DO NOT UNCOMMENT THIS df1.to_csv("/home/matilda/PycharmProjects/log_level_estimation/log_level_estimation/3_preprocessed_data/learning_scenario_3/top100_repositories.csv")

# def sytactcial_rule_prediction():
path_store_test_res = "../../5_results/sytactic_evaluation/stored_predictions/"
path_store_test_res_dictionary = "../../5_results/sytactic_evaluation/stored_predictions/"
# all_results = defaultdict(dict)
with open(path_store_test_res_dictionary + "dict_results_.pickle", "rb") as file:
    all_results = pickle.load(file)

print(all_results.keys())

for seed in range(2, 30):

    df = pd.read_csv("../../3_preprocessed_data/learning_scenario_3/top100_processed_manual_fix_processed.csv")
    df = df.loc[:, ["log_messages", "len", "valid_messages", "label"]]

    df = df.reset_index()
    df = df.drop(["index"], axis=1)

    print("Currently exploting the classes ", np.unique(df.label))

    load, labels = extract_load(df)
    class_count = df.groupby("label").count()['len']
    label_mapper = {class_count.index[i]: i for i in range(len(class_count))}
    tokenized, labels_tokenized, tokenizer = tokenization_dataset(df, load, labels, label_mapper)
    assert len(tokenized) == df.shape[0], "Some data samples have been lost during tokenization. Take care of this."

    from sklearn.model_selection import StratifiedShuffleSplit

    # seed = 10
    stSampling = StratifiedShuffleSplit(1, test_size=0.3, random_state=seed)

    load_train1 = np.array(tokenized, dtype=object)
    load_train_labels1 = np.array(labels_tokenized)

    df1 = copy.copy(df)
    for train_index, test_index in stSampling.split(load_train1, load_train_labels1):

        learning_rate = 0.0001
        device = "gpu"
        decay = 0.001
        betas = (0.8, 0.999)
        momentum = 0.8
        n_epochs = 200
        batch_size = 64
        pad_len = max_len = 32
        n_layers = 8
        in_features = 16
        out_features = 16
        num_heads = 8
        dropout = 0.05

        n_classes = 2
        random_seed = seed
        torch.manual_seed(random_seed)
        np.random.seed(random_seed)

        load_train = load_train1[train_index]
        load_train_labels = load_train_labels1[train_index]
        load_test = load_train1[test_index]
        load_test_labels = load_train_labels1[test_index]

        train_dataloader, test_dataloader = create_data_loaders(load_train, load_train_labels, load_test,  load_test_labels, pad_len, batch_size)

        if device == "gpu":
            torch.cuda.empty_cache()

        src_vocab = tokenizer.n_words

        calculate_weights = lambda x, i: x.sum() / (len(x)*x[i])
        weights = [calculate_weights(class_count,i) for i in range(len(class_count))]

        # weights /= max(weights)

        weights = np.array([1., 1]) # weights = np.array([2.5, 1])
        # weights = max(weights) / weights
        # tmp = weights[1]
        # weights[1] = weights[0]
        # weights[0] = tmp

        if device == "gpu":
            class_weights=torch.FloatTensor(weights).cuda()
            cross_entropoy_loss = nn.CrossEntropyLoss(weight=class_weights).cuda()
        else:
            class_weights = torch.FloatTensor(weights)
            cross_entropoy_loss = nn.CrossEntropyLoss(weight=class_weights)

        loss_f = cross_entropoy_loss
        #
        model = NuLogsyModel(src_vocab=src_vocab, tgt_vocab=n_classes,
                             n_layers=n_layers, in_features=in_features,
                             out_features=out_features,num_heads=num_heads,
                             dropout=dropout, max_len=max_len).get_model()

        if device == "gpu":
            torch.cuda.set_device(0)
            model.cuda()

        sgd_opt = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=decay)
        adam_opt = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=betas, weight_decay=decay)
        optimizers = {"adam":adam_opt,"sgd":sgd_opt}
        optimizer = optimizers['adam']

        model, preds, conf_matrix, best_f1, epoch_idx = run_optimizer(model, train_dataloader, test_dataloader, load_test_labels, optimizer, n_epochs, cross_entropoy_loss,polars=None,class_weights=weights, device=device)

        tokenized, labels_tokenized, tokenizer = word2_vec_representation(df, load, labels, label_mapper, nlp)
        assert len(tokenized) == df.shape[0], "Some data samples have been lost during tokenization. Take care of this."

        load_train = np.array(tokenized, dtype=np.float32)[train_index]
        load_train_labels = np.array(labels_tokenized)[train_index]

        load_test_good_repos = np.array(tokenized, dtype=np.float32)[test_index]
        load_test_good_repos_labels = np.array(labels_tokenized)[test_index]

        print("STARTED TRAINING RANDOM FOREST!")

        model2 = make_pipeline(StandardScaler(), RandomForestClassifier(n_estimators=100, max_depth=None, max_features=int(tokenized[0].shape[0]/ 3)))
        model2.fit(load_train, load_train_labels)
        preds_good_m2 = model2.predict(load_test_good_repos)

        print("STARTED TRAINING SUPPORT VECTOR MACHINE!")
        model3 = make_pipeline(StandardScaler(), SVC(gamma="auto", kernel="rbf"))
        model3.fit(load_train, load_train_labels)
        preds_good_m3 = model3.predict(load_test_good_repos)

        df_train = df.loc[train_index]
        df_test = df.loc[test_index]

        load_test_labels, synct_pred, test_data = syntactic_rule(df_train, df_test, load_test_labels)

        res_seed = defaultdict(dict)

        att = evaluate(preds, load_test_labels)
        att['best_f1'] = best_f1
        att['best_f1_epoch_idx'] = epoch_idx
        att['f1_conf_matricies'] = conf_matrix

        rf = evaluate(preds_good_m2, load_test_labels)
        svm = evaluate(preds_good_m3, load_test_labels)
        synt = evaluate(synct_pred, load_test_labels)

        res_seed["attention"] = att
        res_seed["RandomForest"] = rf
        res_seed["Svm"] = svm
        res_seed["SyntacticRules"] = synt

        all_results[seed] = res_seed

        df_test["labels_attn"] = preds
        df_test["labels_RF"] = preds_good_m2
        df_test["labels_SVM"] = preds_good_m3
        df_test.to_csv(path_store_test_res+"results_"+str(seed), index=False)

        with open(path_store_test_res_dictionary+"dict_results_.pickle", "wb") as file:
            pickle.dump(all_results, file)

# models_location = "../../5_results/models/learning_scenario_2/"
# model_name = models_location + "/" +  scenario + "/info_error_warning_train_1_test_2345.pth"
# tokenizer_name = models_location + scenario + "/info_error_warning_train_1_test_2345_tokenizer.pickle"
# conf_matrix_name = models_location + scenario + "/info_error_warning_train_1_test_2345_confusion_matrix.pickle"
# label_mapper_name = models_location + scenario + "/mapper_train_1_test_2345.pickle"

# torch.save(model, model_name)
# with open(tokenizer_name, "wb") as file:
#     pickle.dump(tokenizer, file)
# with open(conf_matrix_name, "wb") as file:
#     pickle.dump(conf_matrix, file)
# with open(label_mapper_name, "wb") as file:
#     pickle.dump(label_mapper, file)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import argparse

import torch
import torch.nn.parallel
import torch.utils.data
import torchvision.transforms as transforms
import utils
from utils import PointLoss
from utils import distance_squre
import data_utils as d_utils
import ModelNet40Loader
import shapenet_part_loader
from model_FPNet import _netlocalD,_netG

parser = argparse.ArgumentParser()
#parser.add_argument('--dataset',  default='ModelNet40', help='ModelNet10|ModelNet40|ShapeNet')
parser.add_argument('--dataroot',  default='dataset/train', help='path to dataset')
parser.add_argument('--workers', type=int,default=2, help='number of data loading workers')
parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
parser.add_argument('--pnum', type=int, default=2048, help='the point number of a sample')
parser.add_argument('--crop_point_num',type=int,default=512,help='0 means do not use else use with this weight')
parser.add_argument('--nc', type=int, default=3)
parser.add_argument('--niter', type=int, default=300, help='number of epochs to train for')
parser.add_argument('--weight_decay', type=float, default=0.001)
parser.add_argument('--learning_rate', default=0.0002, type=float, help='learning rate in training')
parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. default=0.9')
parser.add_argument('--cuda', type = bool, default = False, help='enables cuda')
parser.add_argument('--ngpu', type=int, default=2, help='number of GPUs to use')
parser.add_argument('--netG', default='Checkpoint/point_netG.pth', help="path to netG (to continue training)")
parser.add_argument('--infile',type = str, default = 'test_one/crop4-1.csv')
parser.add_argument('--infile_real',type = str, default = 'test_one/real4-1.csv')
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--manualSeed', type=int, help='manual seed')
parser.add_argument('--drop',type=float,default=0.2)
parser.add_argument('--num_scales',type=int,default=3,help='number of scales')
# If you want to test your point clouds.  Set the first parameter of '--point_scales_list' equal to (point_number + 512).
parser.add_argument('--point_scales_list',type=list,default=[2048,1024,512],help='number of points in each scales')
parser.add_argument('--each_scales_size',type=int,default=1,help='each scales size')
parser.add_argument('--wtl2',type=float,default=0.9,help='0 means do not use else use with this weight')
parser.add_argument('--cropmethod', default = 'random_center', help = 'random|center|random_center')
opt = parser.parse_args()
print(opt)

def distance_squre1(p1,p2):
    return (p1[0]-p2[0])**2+(p1[1]-p2[1])**2+(p1[2]-p2[2])**2

test_dset = shapenet_part_loader.PartDataset( root='./dataset/shapenetcore_partanno_segmentation_benchmark_v0/',classification=True, class_choice='Car', npoints=opt.pnum, split='test')
test_dataloader = torch.utils.data.DataLoader(test_dset, batch_size=opt.batchSize,
                                         shuffle=True,num_workers = int(opt.workers))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
point_netG = _netG(opt.num_scales,opt.each_scales_size,opt.point_scales_list,opt.crop_point_num)
point_netG = torch.nn.DataParallel(point_netG)
point_netG.to(device)
point_netG.load_state_dict(torch.load(opt.netG,map_location=lambda storage, location: storage)['state_dict'])
point_netG.eval()

input_cropped1 = np.loadtxt(opt.infile,delimiter=',')
input_cropped1 = torch.FloatTensor(input_cropped1)
input_cropped1 = torch.unsqueeze(input_cropped1, 0)
Zeros = torch.zeros(1,512,3)
input_cropped1 = torch.cat((input_cropped1,Zeros),1)

input_cropped2_idx = utils.farthest_point_sample(input_cropped1,opt.point_scales_list[1],RAN = True)
input_cropped2     = utils.index_points(input_cropped1,input_cropped2_idx)
input_cropped3_idx = utils.farthest_point_sample(input_cropped1,opt.point_scales_list[2],RAN = False)
input_cropped3     = utils.index_points(input_cropped1,input_cropped3_idx)
input_cropped4_idx = utils.farthest_point_sample(input_cropped1,256,RAN = True)
input_cropped4     = utils.index_points(input_cropped1,input_cropped4_idx)
input_cropped2 = input_cropped2.to(device)
input_cropped3 = input_cropped3.to(device)
input_cropped  = [input_cropped1,input_cropped2,input_cropped3]

fake_center1,fake_center2,fake=point_netG(input_cropped)
fake = fake.cuda()
fake_center1 = fake_center1.cuda()
fake_center2 = fake_center2.cuda()

input_cropped2 = input_cropped2.cpu()
input_cropped3 = input_cropped3.cpu()
input_cropped4 = input_cropped4.cpu()

np_crop2 = input_cropped2[0].detach().numpy()
np_crop3 = input_cropped3[0].detach().numpy()
np_crop4 = input_cropped4[0].detach().numpy()

real = np.loadtxt(opt.infile_real,delimiter=',')
real = torch.FloatTensor(real)
real = torch.unsqueeze(real,0)
real2_idx = utils.farthest_point_sample(real,64, RAN = False)
real2 = utils.index_points(real,real2_idx)
real3_idx = utils.farthest_point_sample(real,128, RAN = True)
real3 = utils.index_points(real,real3_idx)

real2 = real2.cpu()
real3 = real3.cpu()

np_real2 = real2[0].detach().numpy()
np_real3 = real3[0].detach().numpy()

fake =fake.cpu()
fake_center1 = fake_center1.cpu()
fake_center2 = fake_center2.cpu()
np_fake = fake[0].detach().numpy()
np_fake1 = fake_center1[0].detach().numpy()
np_fake2 = fake_center2[0].detach().numpy()
input_cropped1 = input_cropped1.cpu()
np_crop = input_cropped1[0].numpy()

np.savetxt('test_one/crop_ours'+'.csv', np_crop, fmt = "%f,%f,%f")
np.savetxt('test_one/fake_ours'+'.csv', np_fake, fmt = "%f,%f,%f")
np.savetxt('test_one/crop_ours_txt'+'.txt', np_crop, fmt = "%f,%f,%f")
np.savetxt('test_one/fake_ours_txt'+'.txt', np_fake, fmt = "%f,%f,%f")


import numpy
from gnuradio import gr
import serial
from threading import Timer
from time import sleep

class pwsat2_ptt(gr.sync_block):
    """
    This block is used to control PTT state by serial interface.

    The transciever's PTT input is wired to RTS signal of COM port. Low state ("False") is receive, High state ("True") is transmit.

    The block should be placed right before the audio sink where signal is transmitted to transciever.

    The block works as follow:
    * During initialization it opens serial port and sets RTS signal to Low ("RX")
    * When it gets any samples on input it:
        * swiches RTS signal to High ("TX")
        * optionally waits predelay interval to propagate start of transmission to Transciever
        * calculates time that the whole buffer will take to transmitt
        * copies input to output (it is done to ensure the audio sink will be placed after this block so any COM delay will be less risk)
        * adds delay to that time and creates timer that will shut off TX state.
    * When another batch of samples would arive, the existing timer is cancelled and there is created new timer. This is done to ensure no break during transmission occur.
    * When timer wakes up, the RTS line is switched to False, that stops TX mode
    """
    def __init__(self, device, samp_rate, delay, predelay, is_predelay_used):
        gr.sync_block.__init__(self,
            name="pwsat2_ptt",
            in_sig=[numpy.float32],
            out_sig=[numpy.float32])

        self.ptt = serial.Serial()
        self.device = device
        self.samp_rate = samp_rate
        self.delay = delay
        self.predelay = predelay
        self.timer = None
        self.is_predelay_used = is_predelay_used
        self.is_predelay_fired = False

    def start(self):
        try:
            self.ptt.rts = False
            self.ptt.port = self.device
            self.ptt.open()
            self.is_predelay_fired = False
            return True
        except:
            return False

    def stop(self):
        self.ptt.close()
        return True

    def stop_tx(self):
        self.ptt.rts = False
        self.is_predelay_fired = False

    def work(self, input_items, output_items):
        self.ptt.rts = True
        if self.timer:
            self.timer.cancel()

        in0 = input_items[0]
        out = output_items[0]

        length = len(in0)
        time_of_packet = float(length) / float(self.samp_rate)
        total_time = time_of_packet + self.delay

        out[:] = in0

        if self.is_predelay_used and not self.is_predelay_fired:
            sleep(self.predelay)
            self.is_predelay_fired = True

        self.timer = Timer(total_time, self.stop_tx)
        self.timer.start()

        return len(output_items[0])


'''
Class for experience replay with multiple actors
'''

from __future__ import print_function
import numpy as np
import random
import matplotlib.pyplot as plt
from matplotlib import gridspec
import time
import os
from . import util as my_util

class MultiExperienceMemory:

    def __init__(self, args, multi_simulator = None, target_maker = None):
        ''' Initialize emtpy experience dataset.
            Assumptions:
                - observations come in sequentially, and there is a terminal state in the end of each episode
                - every episode is shorter than the memory
        '''

        # params
        self.capacity = int(args['memory_capacity'])
        self.history_length = int(args['history_length'])
        self.history_step = int(args['history_step'])
        self.shared = args['shared']
        self.obj_shape = args['obj_shape']

        self.num_heads = int(multi_simulator.num_simulators)
        self.target_maker = target_maker
        self.head_offset = int(self.capacity/self.num_heads)

        self.img_shape = (multi_simulator.num_channels, multi_simulator.resolution[1], multi_simulator.resolution[0])
        self.meas_shape = (multi_simulator.num_meas,)
        self.action_shape = (multi_simulator.action_len,)
        self.state_imgs_shape = (self.history_length*self.img_shape[0],) +  self.img_shape[1:]
        self.state_meas_shape = (self.history_length*self.meas_shape[0],)

        # initialize dataset
        self.reset()

    def reset(self):
        self._images = my_util.make_array(shape=(self.capacity,) + self.img_shape, dtype=np.uint8, shared=self.shared, fill_val=0)
        self._measurements =  my_util.make_array(shape=(self.capacity,) + self.meas_shape, dtype=np.float32, shared=self.shared, fill_val=0.)
        self._rewards =  my_util.make_array(shape=(self.capacity,), dtype=np.float32, shared=self.shared, fill_val=0.)
        self._terminals =  my_util.make_array(shape=(self.capacity,), dtype=np.bool, shared=self.shared, fill_val=1)
        self._actions = my_util.make_array(shape=(self.capacity,) + self.action_shape, dtype=np.int, shared=self.shared, fill_val=0)
        self._objectives = my_util.make_array(shape=(self.capacity,) + self.obj_shape, dtype=np.float32, shared=self.shared, fill_val=0.)
        self._n_episode = my_util.make_array(shape=(self.capacity,), dtype=np.uint64, shared=self.shared, fill_val=0) # this is needed to compute future targets efficiently
        self._n_head = my_util.make_array(shape=(self.capacity,), dtype=np.uint64, shared=self.shared, fill_val=0) # this is needed to compute future targets efficiently

        self._curr_indices = np.arange(self.num_heads) * int(self.head_offset)
        self._episode_counts = np.zeros(self.num_heads)

    def add(self, imgs, meass, rwrds, terms, acts, objs=None, preds=None):
        ''' Add experience to dataset.

        Args:
            img: single observation frame
            meas: extra measurements from the state
            rwrd: reward
            term: terminal state
            act: action taken
        '''

        self._images[self._curr_indices] = imgs
        self._measurements[self._curr_indices] = meass
        self._rewards[self._curr_indices] = rwrds
        self._terminals[self._curr_indices] = terms
        self._actions[self._curr_indices] = np.array(acts)
        if isinstance(objs, np.ndarray):
            self._objectives[self._curr_indices] = objs
        if isinstance(preds, np.ndarray):
            self._predictions[self._curr_indices] = preds
        self._n_episode[self._curr_indices] = self._episode_counts
        # this is a hack to simulate our version of ViZDoom which gives the agent the first post-mortem measurement. This turns out to matter a bit for learning
        terminated = np.where(np.array(terms) == True)[0]
        term_inds = self._curr_indices[terminated]
        self._measurements[term_inds] = self._measurements[(term_inds-1)%self.capacity]
        if self.meas_shape[0] == 1:
            for ti in term_inds:
                if self._measurements[ti,0] < 8.:
                    self._measurements[ti,0] -= 8.
        if self.meas_shape[0] == 3:
            for ti in term_inds:
                if self._measurements[ti,1] < 12.:
                    self._measurements[ti,1] -= 12.
        # in case there are 2 terminals in a row - not sure this actually can happen, but just in case
        prev_terminals = np.where(self._terminals[(term_inds-1)%self.capacity])[0]
        if len(prev_terminals) > 0:
            print('Prev terminals', prev_terminals)
        self._measurements[term_inds[prev_terminals]] = 0.
        # end of hack
        # self._n_episode[term_inds] = self._episode_counts[terminated]+100# so that the terminal step of the episode is not used as a target by target_maker
        self._n_head[self._curr_indices] = np.arange(self.num_heads)

        self._episode_counts = self._episode_counts + (np.array(terms) == True)
        self._curr_indices = (self._curr_indices + 1) % self.capacity
        self._terminals[self._curr_indices] = True # make the following state terminal, so that our current episode doesn't get stitched with the next one when sampling states

    def add_step(self, multi_simulator, acts = None, objs=None, preds=None):
        if acts == None:
            acts = multi_simulator.get_random_actions()
        self.add(*(multi_simulator.step(acts) +  (acts,objs,preds)))

    def add_n_steps_with_actor(self, multi_simulator, num_steps, actor, verbose=False, write_predictions=False, write_logs = False, global_step=0):
        ns = 0
        last_meas = np.zeros((multi_simulator.num_simulators,) + self.meas_shape)
        if write_predictions and not hasattr(self,'_predictions'):
            self._predictions = my_util.make_array(shape=(self.capacity,) + actor.predictions_shape, dtype=np.float32, shared=self.shared, fill_val=0.)
        if verbose or write_logs:
            start_time = time.time()
        #write_logs = False
        if write_logs:
            log_dir = os.path.dirname(self.log_prefix)
            if log_dir and not os.path.exists(log_dir):
                os.makedirs(log_dir)
            log_brief = open(self.log_prefix + '_brief.txt','a')
            log_detailed = open(self.log_prefix + '_detailed.txt','a')
            log_detailed.write('Step {0}\n'.format(global_step))
            start_times = time.time() * np.ones(multi_simulator.num_simulators)
            num_episode_steps = np.zeros(multi_simulator.num_simulators)
            accum_rewards = np.zeros(multi_simulator.num_simulators)
            accum_meas = np.zeros((multi_simulator.num_simulators,) + self.meas_shape)
            total_final_meas = np.zeros(self.meas_shape)
            total_avg_meas = np.zeros(self.meas_shape)
            total_accum_reward = 0
            total_start_time = time.time()
            num_episodes = 0
            meas_dim = np.prod(self.meas_shape)
            log_brief_format = ' '.join([('{' + str(n) + '}') for n in range(5)]) + ' | ' + \
                       ' '.join([('{' + str(n+5) + '}') for n in range(meas_dim)]) + ' | ' + \
                       ' '.join([('{' + str(n+5+meas_dim) + '}') for n in range(meas_dim)]) + '\n'
            log_detailed_format = ' '.join([('{' + str(n) + '}') for n in range(4)]) + ' | ' + \
                          ' '.join([('{' + str(n+4) + '}') for n in range(meas_dim)]) + ' | ' + \
                          ' '.join([('{' + str(n+4+meas_dim) + '}') for n in range(meas_dim)]) + '\n'
        for ns in range(int(num_steps)):
            if verbose and time.time() - start_time > 2:
                print('%d/%d' % (ns * multi_simulator.num_simulators, num_steps * multi_simulator.num_simulators))
                start_time = time.time()

            curr_act = actor.act_with_multi_memory(self)

            # actor has to return a np array of bools
            invalid_states = np.logical_not(np.array(self.curr_states_with_valid_history()))
            if actor.random_objective_coeffs:
                actor.reset_objective_coeffs(np.where(invalid_states)[0].tolist())
            curr_act[invalid_states] = actor.random_actions(np.sum(invalid_states))

            if write_predictions:
                self.add_step(multi_simulator, curr_act.tolist(), actor.objectives_to_write(), actor.curr_predictions)
            else:
                self.add_step(multi_simulator, curr_act.tolist(), actor.objectives_to_write())
            if write_logs:
                last_indices = np.array(self.get_last_indices())
                last_rewards = self._rewards[last_indices]
                prev_meas = last_meas
                last_meas = self._measurements[last_indices]
                last_terminals = self._terminals[last_indices]
                last_meas[np.where(last_terminals)[0]] = 0
                accum_rewards += last_rewards
                accum_meas += last_meas
                num_episode_steps = num_episode_steps + 1
                terminated_simulators = list(np.where(last_terminals)[0])
                for ns in terminated_simulators:
                    num_episodes += 1
                    episode_time = time.time() - start_times[ns]
                    avg_meas = accum_meas[ns]/float(num_episode_steps[ns])
                    total_avg_meas += avg_meas
                    total_final_meas += prev_meas[ns]
                    total_accum_reward += accum_rewards[ns]
                    start_times[ns] = time.time()
                    log_detailed.write(log_detailed_format.format(*([num_episodes, num_episode_steps[ns], episode_time, accum_rewards[ns]] + list(prev_meas[ns]) + list(avg_meas))))
                    accum_meas[ns] = 0
                    accum_rewards[ns] = 0
                    num_episode_steps[ns] = 0
                    start_times[ns] = time.time()
        if write_logs:
            if num_episodes == 0:
                num_episodes = 1
            log_brief.write(log_brief_format.format(*([global_step, time.time(), time.time() - total_start_time, num_episodes, total_accum_reward/float(num_episodes)] +
                                 list(total_final_meas / float(num_episodes)) + list(total_avg_meas / float(num_episodes)))))
            log_brief.close()
            log_detailed.close()

    def get_states(self, indices):

        frames = np.zeros(len(indices)*self.history_length, dtype=np.int64)
        for (ni,index) in enumerate(indices):
            frame_slice = np.arange(int(index) - self.history_length*self.history_step + 1, (int(index) + 1), self.history_step) % self.capacity
            frames[ni*self.history_length:(ni+1)*self.history_length] = frame_slice

        state_imgs = np.transpose(np.reshape(np.take(self._images, frames, axis=0), (len(indices),) + self.state_imgs_shape), [0,2,3,1]).astype(np.float32)
        state_meas = np.reshape(np.take(self._measurements, frames, axis=0), (len(indices),) + self.state_meas_shape).astype(np.float32)
        return state_imgs, state_meas

    def get_current_state(self):
        '''  Return most recent observation sequence '''
        return self.get_states(list((self._curr_indices-1)%self.capacity))

    def get_last_indices(self):
        '''  Return most recent indices '''
        return list((self._curr_indices-1)%self.capacity)

    def get_targets(self, indices):
        # TODO this 12345678 is a hack, but should be good enough
        return self.target_maker.make_targets(indices, self._measurements, self._rewards, self._n_episode + 12345678*self._n_head)

    def has_valid_history(self, index):
        return (not self._terminals[np.arange(int(index) - self.history_length*self.history_step + 1, int(index)+1) % self.capacity].any())

    def curr_states_with_valid_history(self):
        return [self.has_valid_history((ind - 1)%self.capacity) for ind in list(self._curr_indices)]

    def has_valid_target(self, index):
        return (not self._terminals[np.arange(index, index+self.target_maker.min_future_frames+1) % self.capacity].any())

    def is_valid_state(self, index):
        return self.has_valid_history(index) and self.has_valid_target(index)

    def get_observations(self, indices):
        indices_arr = np.array(indices)
        state_imgs, state_meas = self.get_states((indices_arr - 1) % self.capacity)
        rwrds = self._rewards[indices_arr]
        acts = self._actions[indices_arr]
        terms = self._terminals[indices_arr].astype(int)
        targs = self.get_targets((indices_arr - 1) % self.capacity)
        if isinstance(self._objectives, np.ndarray):
            objs = self._objectives[indices_arr]
        else:
            objs = None

        return state_imgs, state_meas, rwrds, terms, acts, targs, objs

    def get_random_batch(self, batch_size):
        ''' Sample minibatch of experiences for training '''

        samples = [] # indices of the end of each sample

        while len(samples) < batch_size:
            index = random.randrange(self.capacity)
            # check if there is enough history to make a state and enough future to make targets
            if self.is_valid_state(index):
                samples.append(index)
            else:
                continue

        # create batch
        return self.get_observations(np.array(samples))

    def compute_avg_meas_and_rwrd(self, start_idx, end_idx):
        # compute average measurement values per episode, and average cumulative reward per episode
        curr_num_obs = 0.
        curr_sum_meas = self._measurements[0] * 0
        curr_sum_rwrd = self._rewards[0] * 0
        num_episodes = 0.
        total_sum_meas = self._measurements[0] * 0
        total_sum_rwrd = self._rewards[0] * 0
        for index in range(int(start_idx), int(end_idx)):
            curr_sum_rwrd += self._rewards[index]
            if self._terminals[index]:
                if curr_num_obs:
                    total_sum_meas += curr_sum_meas / curr_num_obs
                    total_sum_rwrd += curr_sum_rwrd
                    num_episodes += 1
                curr_sum_meas = self._measurements[0] * 0
                curr_sum_rwrd = self._rewards[0] * 0
                curr_num_obs = 0.
            else:
                curr_sum_meas += self._measurements[index]
                curr_num_obs += 1

        if num_episodes == 0.:
            total_avg_meas = curr_sum_meas / curr_num_obs
            total_avg_rwrd = curr_sum_rwrd
        else:
            total_avg_meas = total_sum_meas / num_episodes
            total_avg_rwrd = total_sum_rwrd / num_episodes

        return total_avg_meas, total_avg_rwrd

    def show(self, start_index=0, end_index=None, display=True, write_imgs=False, write_video = False, preprocess_targets=None, show_predictions=0, net_discrete_actions = []):
        if show_predictions:
            assert(hasattr(self,'_predictions'))
        curr_index = start_index
        if not end_index:
            end_index = start_index
        inp = ''
        if write_imgs:
            os.makedirs('imgs')
            prev_time = time.time()
        if write_video:
            print(self.img_shape)
            vw = VideoWriter('vid.avi', (self.img_shape[2],self.img_shape[1]), framerate=24, rgb=(self.img_shape[0]==3), mode='replace')
        print('Press ENTER to go to the next observation, type "quit" or "q" or "exit" and press ENTER to quit')

        if display or write_imgs:
            fig_img = plt.figure(figsize=(10, 7), dpi=50, tight_layout=True)
            ax_img = plt.gca()
        if display:
            fig_img.show()
        ns = 0
        if show_predictions and len(net_discrete_actions):
            action_labels = []
            for act in net_discrete_actions:
                action_labels.append(''.join(str(int(i)) for i in act))

        while True:
            curr_img = np.transpose(self._images[curr_index], (1,2,0))
            if curr_img.shape[2] == 1:
                curr_img = np.tile(curr_img, (1,1,3))
            if show_predictions:
                preds = self._predictions[curr_index]
                objs = np.sum(preds, axis=1)
                objs_argsort = np.argsort(-objs)
                curr_preds = np.transpose(preds[objs_argsort[:show_predictions]])
                curr_labels = [action_labels[i] for i in objs_argsort[:show_predictions]]

            if curr_index == start_index:
                if display or write_imgs:
                    im = ax_img.imshow(curr_img)
                    txt = ax_img.text(self.img_shape[2] - 10*len(self._measurements[curr_index]) , self.img_shape[1] - 5, str(self._measurements[curr_index]), fontsize=20, color='red')
                if show_predictions:
                    all_objs = np.sum(self._predictions, axis=2)
                    sbp = my_util.StackedBarPlot(curr_preds, ylim=[np.min(all_objs), np.max(all_objs)], labels=curr_labels)
                    del all_objs
                    if display:
                        sbp.show()
            else:
                if display or write_imgs:
                    im.set_data(curr_img)
                    txt.set_text(str(self._measurements[curr_index]))
                if show_predictions:
                    sbp.set_data(curr_preds, labels = curr_labels)
            if write_imgs:
                plt.savefig('imgs/%.5d.png' % curr_index, dpi=50)
                if time.time() - prev_time > 2:
                    print('Wrote %d images' % ns)
                    prev_time = time.time()
            if write_video:
                vw.add_frame(curr_img[:,:,0])
            if display:
                fig_img.canvas.draw()
                if show_predictions:
                    sbp.draw()
                print('Index', curr_index)
                print('Measurements:', self._measurements[curr_index])
                print('Rewards:', self._rewards[curr_index])
                print('Action:', self._actions[curr_index])
                print('Terminal:', self._terminals[curr_index])
                inp = input()

            curr_index = (curr_index + 1) % self.capacity
            if curr_index == end_index:
                if write_video:
                    vw.close()
                break
            if inp == 'q' or inp == 'quit' or inp == 'exit':
                break
            ns += 1


'''
Turn a tfrecord of NoteSequences into an HDF5 dataset with instruments split
and the pretrained trio mode's latent vectors.
'''

from copy import deepcopy
from time import time
import itertools as it

import click
import logging
from pprint import pformat

import numpy as np
import h5py

import tensorflow as tf
tf.logging.set_verbosity(tf.logging.ERROR)  # noqa

import magenta.music as mm
from magenta.models.music_vae import configs
from magenta.models.music_vae.trained_model import TrainedModel

from constants import (TIMESTEPS, DIM_MELODY, DIM_BASS, DIM_DRUMS, DIM_TRIO)

# Constants
MODEL_NAME = 'hierdec-trio_16bar'

@click.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.argument('output_file', type=click.Path(exists=False))
@click.option('--include_all_instruments', type=bool, default=False)
@click.option('--chunk_size', type=click.IntRange(min=1), default=128,
              help='Number of MIDI files to read at once.')
@click.option('--buffer_size', type=click.IntRange(min=1), default=50000,
              help='Number of examples to make room for at a time.')
@click.option('--batch_size', type=click.IntRange(min=1), default=256,
              help='Batch size for the pretrained model.')
@click.option('--checkpoint', type=click.Path(),
              default='./models/pretrained/{}.ckpt'.format(MODEL_NAME),
              help='Checkpoint to use for the pretrained model.')
@click.option('--log_period', type=click.IntRange(min=0), default=1,
              help='How many chunks pass between logging lines.')
@click.option('--log_file', type=click.Path(),
              default='logs/split_dataset.log')
def main(input_file, output_file,
         include_all_instruments, chunk_size, buffer_size, batch_size,
         checkpoint, log_period, log_file):
    args = locals()
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    log = logging.getLogger(__name__)
    handler = logging.FileHandler(log_file)
    handler.setFormatter(formatter)
    handler.setLevel(logging.DEBUG)
    log.addHandler(handler)
    log.setLevel(logging.DEBUG)

    log.info('Generating melody dataset with args:\n' + pformat(args))
    total_start_time = time()
    ns_gen = mm.note_sequence_io.note_sequence_record_iterator(input_file)
    ns_iter = iter(ns_gen)
    config = configs.CONFIG_MAP[MODEL_NAME]
    trio_converter = config.data_converter

    log.debug('Creating HDF5 store...')
    start_time = time()
    with h5py.File(output_file, 'w') as data_file:
        dataset_size = buffer_size
        ds_melody = data_file.create_dataset(
            'melody',
            (dataset_size, TIMESTEPS, DIM_MELODY),
            maxshape=(None, TIMESTEPS, DIM_MELODY),
            dtype=np.bool
        )
        ds_code = data_file.create_dataset(
            'code',
            (dataset_size, config.hparams.z_size),
            maxshape=(None, config.hparams.z_size),
            dtype=np.float32
        )

        if include_all_instruments:
            ds_trio = data_file.create_dataset(
                'trio',
                (dataset_size, TIMESTEPS, DIM_TRIO),
                maxshape=(None, TIMESTEPS, DIM_TRIO),
                dtype=np.bool
            )
            ds_bass = data_file.create_dataset(
                'bass',
                (dataset_size, TIMESTEPS, DIM_BASS),
                maxshape=(None, TIMESTEPS, DIM_BASS),
                dtype=np.bool
            )
            ds_drums = data_file.create_dataset(
                'drums',
                (dataset_size, TIMESTEPS, DIM_DRUMS),
                maxshape=(None, TIMESTEPS, DIM_DRUMS),
                dtype=np.bool
            )

        log.debug('Done creating HDF5 store (time: {0:.1f}s)'
                  .format(time() - start_time))

        log.debug('Loading model...')
        start_time = time()
        model = TrainedModel(config, batch_size=batch_size,
                             checkpoint_dir_or_path=checkpoint)
        log.debug('Done loading model (time: {0:.1f}s)'
                  .format(time() - start_time))

        log.info('Beginning dataset creation...')
        i_chunk = 0
        i_example = 0
        try:
            while True:
                i_chunk += 1
                log.disabled = i_chunk % log_period != 0 or not log_period
                chunk_time = time()

                log.debug('Processing a chunk of NoteSequences...')
                start_time = time()

                note_sequences = list(it.islice(ns_iter, chunk_size))
                if not note_sequences:
                    break

                trio_tensors = map(
                    lambda seq: trio_converter.to_tensors(seq).outputs,
                    note_sequences
                )
                trio_tensors = it.chain.from_iterable(trio_tensors)
                trio_tensors = list(
                    filter(lambda t: t.shape == (TIMESTEPS, DIM_TRIO),
                           trio_tensors)
                )

                # Ensure an example doesn't overflow the allocated space
                trio_tensors = trio_tensors[:buffer_size]
                n_tensors = len(trio_tensors)
                i_last = n_tensors + i_example

                melody_tensors = list(map(lambda t: t[:, :DIM_MELODY],
                                          trio_tensors))

                if include_all_instruments:
                    bass_tensors = list(map(
                        lambda t: t[:, DIM_MELODY:DIM_MELODY + DIM_BASS],
                        trio_tensors
                    ))
                    drums_tensors = list(map(lambda t: t[:, -DIM_DRUMS:],
                                             trio_tensors))

                log.debug('Done processing NoteSequences (time: {0:.1f}s)'
                          .format(time() - start_time))

                log.debug('Running encoder...')
                start_time = time()
                _, codes, _ = model.encode_tensors(deepcopy(trio_tensors),
                                                   [TIMESTEPS] * n_tensors)
                log.debug('Done running encoder (time: {0:.1f}s)'
                          .format(time() - start_time))

                if i_last >= dataset_size:
                    dataset_size += buffer_size
                    log.info('Resizing datasets to size:', dataset_size)
                    ds_melody.resize((dataset_size, TIMESTEPS, DIM_MELODY))
                    ds_code.resize((dataset_size, config.hparams.z_size))

                    if include_all_instruments:
                        ds_trio.resize((dataset_size, TIMESTEPS, DIM_TRIO))
                        ds_bass.resize((dataset_size, TIMESTEPS, DIM_BASS))
                        ds_drums.resize((dataset_size, TIMESTEPS, DIM_DRUMS))

                log.debug('Writing examples to HDF5...')
                start_time = time()
                ds_melody[i_example:i_last, :, :] = np.array(melody_tensors)
                ds_code[i_example:i_last, :] = np.array(codes)

                if include_all_instruments:
                    ds_trio[i_example:i_last, :, :] = np.array(trio_tensors)
                    ds_bass[i_example:i_last, :, :] = np.array(bass_tensors)
                    ds_drums[i_example:i_last, :, :] = np.array(drums_tensors)

                log.debug('Done writing examples to HDF5 (time: {0:.1f}s)'
                          .format(time() - start_time))

                i_example += n_tensors

                log.info(('Chunk {0} wrote {1} examples ' +
                         '(total: {2}; time: {3:.1f}s)')
                         .format(i_chunk, n_tensors, i_example,
                                 time() - chunk_time))
        except StopIteration:
            pass

    log.debug('Finished writing data')
    log.debug('Resizing datasets...')
    dataset_size = i_example
    ds_melody.resize((dataset_size, TIMESTEPS, DIM_MELODY))
    ds_code.resize((dataset_size, config.hparams.z_size))
    if include_all_instruments:
        ds_trio.resize((dataset_size, TIMESTEPS, DIM_TRIO))
        ds_bass.resize((dataset_size, TIMESTEPS, DIM_BASS))
        ds_drums.resize((dataset_size, TIMESTEPS, DIM_DRUMS))
    log.debug('Done resizing datasets...')

    total_time = time() - total_start_time
    log.info('Finished creating HDF5 dataset')
    log.info('Total examples: {}'.format(i_example))
    log.info('Total chunks: {}'.format(i_chunk))
    log.info('Total time: {0:.1f}s'.format(total_time))
    log.info('Done!')

if __name__ == '__main__':
    main()


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

from __future__ import absolute_import
import numpy as np
import time
import math
import tqdm
import paddle.fluid as fluid
import paddlex.utils.logging as logging
from paddlex.utils import seconds_to_hms
import paddlex
from collections import OrderedDict
from .base import BaseAPI

class BaseClassifier(BaseAPI):
    """构建分类器，并实现其训练、评估、预测和模型导出。
    Args:
        model_name (str): 分类器的模型名字，取值范围为['ResNet18',
                          'ResNet34', 'ResNet50', 'ResNet101',
                          'ResNet50_vd', 'ResNet101_vd', 'DarkNet53',
                          'MobileNetV1', 'MobileNetV2', 'Xception41',
                          'Xception65', 'Xception71']。默认为'ResNet50'。
        num_classes (int): 类别数。默认为1000。
    """

    def __init__(self, model_name='ResNet50', num_classes=1000):
        self.init_params = locals()
        super(BaseClassifier, self).__init__('classifier')
        if not hasattr(paddlex.cv.nets, str.lower(model_name)):
            raise Exception(
                "ERROR: There's no model named {}.".format(model_name))
        self.model_name = model_name
        self.labels = None
        self.num_classes = num_classes
        self.fixed_input_shape = None

    def build_net(self, mode='train'):
        if self.fixed_input_shape is not None:
            input_shape = [
                None, 3, self.fixed_input_shape[1], self.fixed_input_shape[0]
            ]
            image = fluid.data(
                dtype='float32', shape=input_shape, name='image')
        else:
            image = fluid.data(
                dtype='float32', shape=[None, 3, None, None], name='image')
        if mode != 'test':
            label = fluid.data(dtype='int64', shape=[None, 1], name='label')
        model = getattr(paddlex.cv.nets, str.lower(self.model_name))
        net_out = model(image, num_classes=self.num_classes)
        softmax_out = fluid.layers.softmax(net_out, use_cudnn=False)
        inputs = OrderedDict([('image', image)])
        outputs = OrderedDict([('predict', softmax_out)])
        if mode == 'test':
            self.interpretation_feats = OrderedDict([('logits', net_out)])
        if mode != 'test':
            cost = fluid.layers.cross_entropy(input=softmax_out, label=label)
            avg_cost = fluid.layers.mean(cost)
            acc1 = fluid.layers.accuracy(input=softmax_out, label=label, k=1)
            k = min(5, self.num_classes)
            acck = fluid.layers.accuracy(input=softmax_out, label=label, k=k)
            if mode == 'train':
                self.optimizer.minimize(avg_cost)
            inputs = OrderedDict([('image', image), ('label', label)])
            outputs = OrderedDict([('loss', avg_cost), ('acc1', acc1),
                                   ('acc{}'.format(k), acck)])
        if mode == 'eval':
            del outputs['loss']
        return inputs, outputs

    def default_optimizer(self, learning_rate, lr_decay_epochs, lr_decay_gamma,
                          num_steps_each_epoch):
        boundaries = [b * num_steps_each_epoch for b in lr_decay_epochs]
        values = [
            learning_rate * (lr_decay_gamma**i)
            for i in range(len(lr_decay_epochs) + 1)
        ]
        lr_decay = fluid.layers.piecewise_decay(
            boundaries=boundaries, values=values)
        optimizer = fluid.optimizer.Momentum(
            lr_decay,
            momentum=0.9,
            regularization=fluid.regularizer.L2Decay(1e-04))
        return optimizer

    def train(self,
              num_epochs,
              train_dataset,
              train_batch_size=64,
              eval_dataset=None,
              save_interval_epochs=1,
              log_interval_steps=2,
              save_dir='output',
              pretrain_weights='IMAGENET',
              optimizer=None,
              learning_rate=0.025,
              lr_decay_epochs=[30, 60, 90],
              lr_decay_gamma=0.1,
              use_vdl=False,
              sensitivities_file=None,
              eval_metric_loss=0.05,
              early_stop=False,
              early_stop_patience=5,
              resume_checkpoint=None):
        """训练。
        Args:
            num_epochs (int): 训练迭代轮数。
            train_dataset (paddlex.datasets): 训练数据读取器。
            train_batch_size (int): 训练数据batch大小。同时作为验证数据batch大小。默认值为64。
            eval_dataset (paddlex.datasets: 验证数据读取器。
            save_interval_epochs (int): 模型保存间隔（单位：迭代轮数）。默认为1。
            log_interval_steps (int): 训练日志输出间隔（单位：迭代步数）。默认为2。
            save_dir (str): 模型保存路径。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
                则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
            optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：
                fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
            learning_rate (float): 默认优化器的初始学习率。默认为0.025。
            lr_decay_epochs (list): 默认优化器的学习率衰减轮数。默认为[30, 60, 90]。
            lr_decay_gamma (float): 默认优化器的学习率衰减率。默认为0.1。
            use_vdl (bool): 是否使用VisualDL进行可视化。默认值为False。
            sensitivities_file (str): 若指定为路径时，则加载路径下敏感度信息进行裁剪；若为字符串'DEFAULT'，
                则自动下载在ImageNet图片数据上获得的敏感度信息进行裁剪；若为None，则不进行裁剪。默认为None。
            eval_metric_loss (float): 可容忍的精度损失。默认为0.05。
            early_stop (bool): 是否使用提前终止训练策略。默认值为False。
            early_stop_patience (int): 当使用提前终止训练策略时，如果验证集精度在`early_stop_patience`个epoch内
                连续下降或持平，则终止训练。默认值为5。
            resume_checkpoint (str): 恢复训练时指定上次训练保存的模型路径。若为None，则不会恢复训练。默认值为None。
        Raises:
            ValueError: 模型从inference model进行加载。
        """
        if not self.trainable:
            raise ValueError("Model is not trainable from load_model method.")
        self.labels = train_dataset.labels
        if optimizer is None:
            num_steps_each_epoch = train_dataset.num_samples // train_batch_size
            optimizer = self.default_optimizer(
                learning_rate=learning_rate,
                lr_decay_epochs=lr_decay_epochs,
                lr_decay_gamma=lr_decay_gamma,
                num_steps_each_epoch=num_steps_each_epoch)
        self.optimizer = optimizer
        # 构建训练、验证、预测网络
        self.build_program()
        # 初始化网络权重
        self.net_initialize(
            startup_prog=fluid.default_startup_program(),
            pretrain_weights=pretrain_weights,
            save_dir=save_dir,
            sensitivities_file=sensitivities_file,
            eval_metric_loss=eval_metric_loss,
            resume_checkpoint=resume_checkpoint)
        # 训练
        self.train_loop(
            num_epochs=num_epochs,
            train_dataset=train_dataset,
            train_batch_size=train_batch_size,
            eval_dataset=eval_dataset,
            save_interval_epochs=save_interval_epochs,
            log_interval_steps=log_interval_steps,
            save_dir=save_dir,
            use_vdl=use_vdl,
            early_stop=early_stop,
            early_stop_patience=early_stop_patience)

    def evaluate(self,
                 eval_dataset,
                 batch_size=1,
                 epoch_id=None,
                 return_details=False):
        """评估。
        Args:
            eval_dataset (paddlex.datasets): 验证数据读取器。
            batch_size (int): 验证数据批大小。默认为1。
            epoch_id (int): 当前评估模型所在的训练轮数。
            return_details (bool): 是否返回详细信息。
        Returns:
          dict: 当return_details为False时，返回dict, 包含关键字：'acc1'、'acc5'，
              分别表示最大值的accuracy、前5个最大值的accuracy。
          tuple (metrics, eval_details): 当return_details为True时，增加返回dict，
              包含关键字：'true_labels'、'pred_scores'，分别代表真实类别id、每个类别的预测得分。
        """
        self.arrange_transforms(
            transforms=eval_dataset.transforms, mode='eval')
        data_generator = eval_dataset.generator(
            batch_size=batch_size, drop_last=False)
        k = min(5, self.num_classes)
        total_steps = math.ceil(eval_dataset.num_samples * 1.0 / batch_size)
        true_labels = list()
        pred_scores = list()
        if not hasattr(self, 'parallel_test_prog'):
            self.parallel_test_prog = fluid.CompiledProgram(
                self.test_prog).with_data_parallel(
                    share_vars_from=self.parallel_train_prog)
        batch_size_each_gpu = self._get_single_card_bs(batch_size)
        logging.info(
            "Start to evaluating(total_samples={}, total_steps={})...".format(
                eval_dataset.num_samples, total_steps))
        for step, data in tqdm.tqdm(
                enumerate(data_generator()), total=total_steps):
            images = np.array([d[0] for d in data]).astype('float32')
            labels = [d[1] for d in data]
            num_samples = images.shape[0]
            if num_samples < batch_size:
                num_pad_samples = batch_size - num_samples
                pad_images = np.tile(images[0:1], (num_pad_samples, 1, 1, 1))
                images = np.concatenate([images, pad_images])
            outputs = self.exe.run(
                self.parallel_test_prog,
                feed={'image': images},
                fetch_list=list(self.test_outputs.values()))
            outputs = [outputs[0][:num_samples]]
            true_labels.extend(labels)
            pred_scores.extend(outputs[0].tolist())
            logging.debug("[EVAL] Epoch={}, Step={}/{}".format(
                epoch_id, step + 1, total_steps))

        pred_top1_label = np.argsort(pred_scores)[:, -1]
        pred_topk_label = np.argsort(pred_scores)[:, -k:]
        acc1 = sum(pred_top1_label == true_labels) / len(true_labels)
        acck = sum(
            [np.isin(x, y)
             for x, y in zip(true_labels, pred_topk_label)]) / len(true_labels)
        metrics = OrderedDict([('acc1', acc1), ('acc{}'.format(k), acck)])
        if return_details:
            eval_details = {
                'true_labels': true_labels,
                'pred_scores': pred_scores
            }
            return metrics, eval_details
        return metrics

    def predict(self, img_file, transforms=None, topk=1):
        """预测。
        Args:
            img_file (str): 预测图像路径。
            transforms (paddlex.cls.transforms): 数据预处理操作。
            topk (int): 预测时前k个最大值。
        Returns:
            list: 其中元素均为字典。字典的关键字为'category_id'、'category'、'score'，
            分别对应预测类别id、预测类别标签、预测得分。
        """
        if transforms is None and not hasattr(self, 'test_transforms'):
            raise Exception("transforms need to be defined, now is None.")
        true_topk = min(self.num_classes, topk)
        if transforms is not None:
            self.arrange_transforms(transforms=transforms, mode='test')
            im = transforms(img_file)
        else:
            self.arrange_transforms(
                transforms=self.test_transforms, mode='test')
            im = self.test_transforms(img_file)
        result = self.exe.run(
            self.test_prog,
            feed={'image': im},
            fetch_list=list(self.test_outputs.values()))
        pred_label = np.argsort(result[0][0])[::-1][:true_topk]
        res = [{
            'category_id': l,
            'category': self.labels[l],
            'score': result[0][0][l]
        } for l in pred_label]
        return res

class ResNet18(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet18, self).__init__(
            model_name='ResNet18', num_classes=num_classes)

class ResNet34(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet34, self).__init__(
            model_name='ResNet34', num_classes=num_classes)

class ResNet50(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet50, self).__init__(
            model_name='ResNet50', num_classes=num_classes)

class ResNet101(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet101, self).__init__(
            model_name='ResNet101', num_classes=num_classes)

class ResNet50_vd(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet50_vd, self).__init__(
            model_name='ResNet50_vd', num_classes=num_classes)

class ResNet101_vd(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet101_vd, self).__init__(
            model_name='ResNet101_vd', num_classes=num_classes)

class ResNet50_vd_ssld(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet50_vd_ssld, self).__init__(
            model_name='ResNet50_vd_ssld', num_classes=num_classes)

class ResNet101_vd_ssld(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ResNet101_vd_ssld, self).__init__(
            model_name='ResNet101_vd_ssld', num_classes=num_classes)

class DarkNet53(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(DarkNet53, self).__init__(
            model_name='DarkNet53', num_classes=num_classes)

class MobileNetV1(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(MobileNetV1, self).__init__(
            model_name='MobileNetV1', num_classes=num_classes)

class MobileNetV2(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(MobileNetV2, self).__init__(
            model_name='MobileNetV2', num_classes=num_classes)

class MobileNetV3_small(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(MobileNetV3_small, self).__init__(
            model_name='MobileNetV3_small', num_classes=num_classes)

class MobileNetV3_large(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(MobileNetV3_large, self).__init__(
            model_name='MobileNetV3_large', num_classes=num_classes)

class MobileNetV3_small_ssld(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(MobileNetV3_small_ssld, self).__init__(
            model_name='MobileNetV3_small_ssld', num_classes=num_classes)

class MobileNetV3_large_ssld(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(MobileNetV3_large_ssld, self).__init__(
            model_name='MobileNetV3_large_ssld', num_classes=num_classes)

class Xception65(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(Xception65, self).__init__(
            model_name='Xception65', num_classes=num_classes)

class Xception41(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(Xception41, self).__init__(
            model_name='Xception41', num_classes=num_classes)

class DenseNet121(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(DenseNet121, self).__init__(
            model_name='DenseNet121', num_classes=num_classes)

class DenseNet161(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(DenseNet161, self).__init__(
            model_name='DenseNet161', num_classes=num_classes)

class DenseNet201(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(DenseNet201, self).__init__(
            model_name='DenseNet201', num_classes=num_classes)

class ShuffleNetV2(BaseClassifier):
    def __init__(self, num_classes=1000):
        super(ShuffleNetV2, self).__init__(
            model_name='ShuffleNetV2', num_classes=num_classes)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os, argparse
import numpy as np
import xml.etree.ElementTree as ET
from tqdm import tqdm
from collections import OrderedDict

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test'), ('2012', 'train'), ('2012', 'val')]
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

class_count = {}

def convert_to_darknet(size, box):
    '''
    convert box format from (xmin,xmax,ymin,ymax) to darknet
    style (x,y,w,h) format with relative coordinate
    '''
    dw = 1./size[0]
    dh = 1./size[1]
    x = (box[0] + box[1])/2.0
    y = (box[2] + box[3])/2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(dataset_path, year, image_id, output_path, include_difficult):
    in_file = open('%s/VOC%s/Annotations/%s.xml'%(dataset_path, year, image_id))
    out_file = open('%s/%s.txt'%(output_path, image_id), 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    # get image size (width, height)
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult')
        if difficult is None:
            difficult = '0'
        else:
            difficult = difficult.text
        class_name = obj.find('name').text
        if class_name not in classes:
            continue
        if not include_difficult and int(difficult)==1:
            continue
        class_id = classes.index(class_name)

        # parse box coordinate to (xmin,xmax,ymin,ymax) format
        xml_box = obj.find('bndbox')
        box = (float(xml_box.find('xmin').text), float(xml_box.find('xmax').text), float(xml_box.find('ymin').text), float(xml_box.find('ymax').text))

        # convert box to darknet format and save to txt
        darknet_box = convert_to_darknet((width, height), box)
        out_file.write(str(class_id) + " " + " ".join([str(item) for item in darknet_box]) + '\n')
        class_count[class_name] = class_count[class_name] + 1
    out_file.close()

def has_object(dataset_path, year, image_id, include_difficult):
    '''
    check if an image annotation has valid object bbox info,
    return a boolean result
    '''
    try:
        in_file = open('%s/VOC%s/Annotations/%s.xml'%(dataset_path, year, image_id))
    except:
        # bypass image if no annotation
        return False
    tree=ET.parse(in_file)
    root = tree.getroot()
    count = 0

    for obj in root.iter('object'):
        difficult = obj.find('difficult')
        if difficult is None:
            difficult = '0'
        else:
            difficult = difficult.text
        cls = obj.find('name').text
        if cls not in classes:
            continue
        if not include_difficult and int(difficult)==1:
            continue
        count = count + 1
    return count != 0

def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        classes = f.readlines()
    classes = [c.strip() for c in classes]
    return classes

parser = argparse.ArgumentParser(description='convert PascalVOC dataset annotation to darknet/YOLOv5 txt annotation files')
parser.add_argument('--dataset_path', type=str, help='path to PascalVOC dataset, default=%(default)s', default=os.getcwd()+'/../../VOCdevkit')
parser.add_argument('--year', type=str, help='subset path of year (2007/2012), default will cover both', default=None)
parser.add_argument('--set', type=str, help='convert data set, default will cover train, val and test', default=None)
parser.add_argument('--output_path', type=str,  help='output path for generated annotation txt files, default=%(default)s', default='./')
parser.add_argument('--classes_path', type=str, required=False, help='path to class definitions')
parser.add_argument('--include_difficult', action="store_true", help='to include difficult object', default=False)
parser.add_argument('--include_no_obj', action="store_true", help='to include no object image', default=False)
args = parser.parse_args()

# update class names
if args.classes_path:
    classes = get_classes(args.classes_path)

# get real path for dataset
dataset_realpath = os.path.realpath(args.dataset_path)

# create output path
os.makedirs(args.output_path, exist_ok=True)

# get specific sets to convert
if args.year is not None:
    sets = [item for item in sets if item[0] == args.year]
if args.set is not None:
    sets = [item for item in sets if item[1] == args.set]

for year, image_set in sets:
    # count class item number in each set
    class_count = OrderedDict([(item, 0) for item in classes])

    image_ids = open('%s/VOC%s/ImageSets/Main/%s.txt'%(dataset_realpath, year, image_set)).read().strip().split()
    pbar = tqdm(total=len(image_ids), desc='Converting VOC%s_%s'%(year, image_set))
    for image_id in image_ids:
        file_string = '%s/VOC%s/JPEGImages/%s.jpg'%(dataset_realpath, year, image_id)
        # check if the image file exists
        if not os.path.exists(file_string):
            file_string = '%s/VOC%s/JPEGImages/%s.jpeg'%(dataset_realpath, year, image_id)
        if not os.path.exists(file_string):
            raise ValueError('image file for id: {} not exists'.format(image_id))

        if has_object(dataset_realpath, year, image_id, args.include_difficult):
            convert_annotation(dataset_realpath, year, image_id, args.output_path, args.include_difficult)
        elif args.include_no_obj:
            # include no object image. just create an empty file
            os.system('touch %s/%s.txt'%(args.output_path, image_id))
        pbar.update(1)
    pbar.close()
    # print out item number statistic
    print('\nDone for VOC%s %s. classes number statistic'%(year, image_set))
    print('Image number: %d'%(len(image_ids)))
    print('Object class number:')
    for (class_name, number) in class_count.items():
        print('%s: %d' % (class_name, number))
    print('total object number:', np.sum(list(class_count.values())))