import argparse
import os
import time
import copy
import csv

import numpy as np

from torch import nn
from torch import optim
import torch

from sim_framework.envs.drone_env import DroneEnv

from common import utils
from networks.structures import PolicyNetwork, ValueNetwork, SoftQNetwork

def argparser():

    parser = argparse.ArgumentParser("Benching envs")
    parser.add_argument(
        '--net_size_policy',
        help='Size of the neural networks hidden layers',
        type=int,
        default=64)
    parser.add_argument(
        '--activation_function',
        help='activation function for policy',
        type=str,
        default='relu',
        choices=[
            'relu',
            'tanh'])

    parser.add_argument(
        '--net_size_value',
        help='Size of the neural networks hidden layers',
        type=int,
        default=256)
    parser.add_argument(
        '--max_episodes',
        help='Number of epochs in the training',
        type=int,
        default=int(10000))
    parser.add_argument(
        '--replay_buffer_size',
        help='Size of the replay buffer',
        type=int,
        default=int(1e6))
    parser.add_argument(
        '--num_steps_until_train',
        help='How many steps we sample with current policy',
        type=int,
        default=1)
    parser.add_argument(
        '--num_trains_per_step',
        help='Number of timesteps in each step',
        type=int,
        default=1)
    parser.add_argument(
        '--min_num_steps_before_training',
        help='Number of timesteps using random (uniform) policy to fill \
    the Replay buffer in the beggining of the training',
        type=int,
        default=100)
    parser.add_argument(
        '--batch-size', help='Batch size in each epoch', type=int, default=256)
    # parser.add_argument(
    #     '--use_automatic_entropy_tuning', help='Set True to automatically discover best alpha (temperature)', type=bool, default=True)
    parser.add_argument(
        '--env_reset_mode',
        help='How to sample the starting position of the agent',
        choices=(
            'Uniform',
            'Gaussian',
            'False',
            'Discretized_Uniform'),
        default='False')
    parser.add_argument(
        '--use_cuda',
        help='If the device is a GPU or CPU',
        default='False',
        choices=[
            'True',
            'False'])
    parser.add_argument(
        '--restore_path',
        help='Filename of the policy being loaded',
        type=str,
        default=None)
    parser.add_argument(
        '--save_path', help='Filename to save the current training', type=str)
    parser.add_argument(
        '--log_interval', help='Frequency of logging', type=int, default=100)
    parser.add_argument(
        '--save-interval', help='Frequency of saving', type=int, default=100)
    parser.add_argument(
        '--eval_interval',
        help='Frequency for evaluating deterministic policy',
        type=int,
        default=None)
    parser.add_argument(
        '--use_double', help='Flag to use float64', type=str, default=None)
    parser.add_argument(
        '--learning_rate', help='Learning rate', type=float, default=3e-4)
    parser.add_argument(
        '--reward_function',
        help='What reward function to use',
        default='Normal',
        type=str)
    parser.add_argument(
        '--seed', help='Global seed', default=42, type=int)
    parser.add_argument(
        '--state', help='Global seed', default='Old', type=str)
    parser.add_argument(
        '--same-norm', help='same_norm', default=False)
    parser.add_argument(
        '--threshold',
        help='Clipping the difference between action vectors',
        default=4.0,
        type=float)
    parser.add_argument(
        '--clip-action',
        help='Clipping the difference between action vectors',
        default=100,
        type=int)
    parser.add_argument(
        '--save_interval', help='Frequency of saving', type=int, default=100)
    return parser.parse_args()

    if (args.env_reset_mode) == 'False':
        args.env_reset_mode = False

    if (args.use_cuda) == 'False':
        args.use_cuda = False
    else:
        args.use_cuda = True
    if args.activation_function == 'relu':
        args.activation_function = F.relu
    else:
        args.activation_function = F.tanh

class SAC():

    def __init__(
        self,
        env,
        replay_buffer_size,
        hidden_dim,
        restore_path,
        device,
        max_episodes,
        save_path,
        learning_rate=3e-4,
        use_double=True,
        min_num_steps_before_training=0,
        save_interval=100,
    ):
        """
        Soft Actor-Critic algorithm

        Parameters
        ----------
        env :
            The environment to be used
        replay_buffer_size : [int]
            Replay-buffer size
        hidden_dim : [int]
            Size of the hidden-layers in Q and V functions[description]
        restore_path : [str]
            File path to restore training
        device : [str or torch.device]
            'cpu' or 'gpu
        max_episodes : [int]
            Max number of episodes to train the policy
        save_path : [str]
            File path to save the networks
        learning_rate : [float], optional
            The learning rate for gradient based optimization, by default 3e-4
        use_double : bool, optional
            Use Float Tensor or Double Tensor, by default True
        min_num_steps_before_training : int, optional
            Number of steps to randomly sample before start acting, by default 0
        save_interval : int, optional
            The interval in epochs to save the models, by default 100
        """
        self.env = env

        self.save_path = save_path
        self.save_interval = save_interval
        self.min_num_steps_before_training = min_num_steps_before_training
        self.restore_path = restore_path
        self.device = device
        self.hidden_dim = hidden_dim
        self.use_double = use_double
        # Network and env parameters
        self.action_dim = self.env.action_space.shape[0]
        try:
            self.state_dim = self.env.observation_space.shape[0]
        except BaseException:
            self.state_dim = self.env.observation_space
        # hidden_dim = args.net_size_value
        self.action_range = [self.env.agent.action_space.low.min(
        ), self.env.agent.action_space.high.max()]

        self._creating_models(replay_buffer_size, self.state_dim,
                              self.action_dim, self.device, self.hidden_dim)

        # Copying the data to the target networks
        for target_param, param in zip(
                self.target_value_net.parameters(), self.value_net.parameters()):
            target_param.data.copy_(param.data)

        # Types of losses
        self.value_criterion = nn.MSELoss()
        self.soft_q_criterion = nn.MSELoss()

        # Learning rates
        self.value_lr = learning_rate
        self.soft_q_lr = learning_rate
        self.policy_lr = learning_rate
        # Optimizers
        self.value_optimizer = optim.Adam(
            self.value_net.parameters(), lr=self.value_lr)
        self.soft_q_optimizer = optim.Adam(
            self.soft_q_net.parameters(), lr=self.soft_q_lr)
        self.policy_optimizer = optim.Adam(
            self.policy_net.parameters(), lr=self.policy_lr)

        if self.use_double:
            self.state_to_tensor = lambda x: torch.DoubleTensor(
                x).unsqueeze(0).to(device)
        else:
            self.state_to_tensor = lambda x: torch.FloatTensor(
                x).unsqueeze(0).to(device)

    def soft_q_update(self,
                      batch_size,
                      gamma=0.99,
                      mean_lambda=1e-3,
                      std_lambda=1e-3,
                      z_lambda=0.0,
                      soft_tau=1e-2,
                      ):
        """
        SAC train update (Soft-Q)

        Parameters
        ----------
        batch_size : [int]
            Batch size
        gamma : float, optional
            Discount factor, by default 0.99
        mean_lambda : [float], optional
            coefficient for penalty on policy mean magnitude, by default 1e-3
        std_lambda : [float], optional
            coefficient for penalty on policy variance, by default 1e-3
        z_lambda : float, optional
            coefficient for penalty on policy mean before been squashed by tanh, by default 0
        soft_tau : [float], optional
            Soft coefficient to update target networks, by default 1e-2
        """
        # Sampling memmory batch
        state, action, reward, next_state, done = self.replay_buffer.sample(
            batch_size)

        # Broadcast
        if self.use_double:
            state = torch.DoubleTensor(state).to(self.device)
            next_state = torch.DoubleTensor(next_state).to(self.device)
            action = torch.DoubleTensor(action).to(self.device)
            reward = torch.DoubleTensor(reward).unsqueeze(1).to(self.device)
            done = torch.DoubleTensor(np.float64(
                done)).unsqueeze(1).to(self.device)

        else:
            state = torch.FloatTensor(state).to(self.device)
            next_state = torch.FloatTensor(next_state).to(self.device)
            action = torch.FloatTensor(action).to(self.device)
            reward = torch.FloatTensor(reward).unsqueeze(1).to(self.device)
            done = torch.FloatTensor(np.float32(
                done)).unsqueeze(1).to(self.device)

        # Net forward-passes
        expected_q_value = self.soft_q_net(state, action)
        expected_value = self.value_net(state)
        new_action, log_prob, z, mean, log_std = self.policy_net.evaluate(
            state)

        ## Qf - loss
        target_value = self.target_value_net(next_state)
        next_q_value = reward + (1 - done) * gamma * target_value
        q_value_loss = self.soft_q_criterion(
            expected_q_value, next_q_value.detach())

        ## Vf - loss
        expected_new_q_value = self.soft_q_net(state, new_action)
        next_value = expected_new_q_value - log_prob
        value_loss = self.value_criterion(expected_value, next_value.detach())

        # Policy Loss
        log_prob_target = expected_new_q_value - expected_value
        policy_loss = (log_prob * (log_prob - log_prob_target).detach()).mean()
        mean_loss = mean_lambda * mean.pow(2).mean()
        std_loss = std_lambda * log_std.pow(2).mean()
        z_loss = z_lambda * z.pow(2).sum(1).mean()
        policy_loss += mean_loss + std_loss + z_loss

        # NN updates
        self.soft_q_optimizer.zero_grad()
        q_value_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.value_net.parameters(), 0.5)
        self.soft_q_optimizer.step()

        self.value_optimizer.zero_grad()
        value_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.value_net.parameters(), 0.5)
        self.value_optimizer.step()

        self.policy_optimizer.zero_grad()
        policy_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.value_net.parameters(), 0.5)
        self.policy_optimizer.step()

        # Updating the target networks
        for target_param, param in zip(
                self.target_value_net.parameters(), self.value_net.parameters()):
            target_param.data.copy_(
                target_param.data * (1.0 - soft_tau) + param.data * soft_tau
            )

        return(q_value_loss.item(), policy_loss.item(), value_loss.item())

    def __write_csv(
            self,
            episode,
            time_elapsed,
            frame_count,
            len_replay_buffer,
            episode_reward,
            value_loss,
            q_value_loss,
            policy_loss,
            step):
        """
        Writes data in csv
        """

        with open(os.path.join(self.save_path, 'progress.csv'), 'a') as csvfile:
            rew_writer = csv.writer(csvfile, delimiter=';',
                                    quotechar='|', quoting=csv.QUOTE_MINIMAL)

            rew_writer.writerow([episode,
                                 time_elapsed,
                                 frame_count,
                                 len(self.replay_buffer),
                                 episode_reward,
                                 value_loss,
                                 q_value_loss,
                                 policy_loss,
                                 step])

    def __save_model(self, episode, frame_count):
        """
        Saves model pickles

        Parameters
        ----------
        episode : [int]
            Current episode
        frame_count : [int]
            Current timestep
        """
        save_state = {
            'episode': episode,
            'frame_count': frame_count,
            'value_net': self.value_net.state_dict(),
            'target_value_net': self.target_value_net.state_dict(),
            'soft_q_net': self.soft_q_net.state_dict(),
            'policy_net': self.policy_net.state_dict(),
            'value_optimizer': self.value_optimizer.state_dict(),
            'soft_q_optimizer': self.soft_q_optimizer.state_dict(),
            'policy_optimizer': self.policy_optimizer.state_dict(),
            'replay_buffer': self.replay_buffer
        }
        torch.save(save_state, self.save_path + '/state.pt')
        torch.save(save_state['policy_net'],
                   self.save_path + '/state.pt'[:-3] + '_policy.pt')
        print('saving model at = ', self.save_path)

    def _creating_models(
            self,
            buffer_size,
            state_dim,
            action_dim,
            device,
            hidden_dim):
        """
        Istantiating the networks and buffer
        """

        self.policy_net = PolicyNetwork(
            state_dim,
            action_dim,
            args.net_size_policy).to(device).type(
            torch.double)
        self.eval_policy = copy.deepcopy(self.policy_net)
        self.replay_buffer = utils.ReplayBuffer(buffer_size)
        self.value_net = ValueNetwork(
            state_dim, hidden_dim).to(device).type(torch.double)
        self.target_value_net = ValueNetwork(
            state_dim, hidden_dim).to(device).type(torch.double)
        self.soft_q_net = SoftQNetwork(
            state_dim, action_dim, hidden_dim).to(device).type(torch.double)

        if self.use_double:
            self.policy_net = self.policy_net.type(torch.double)
            self.target_value_net = self.target_value_net.type(torch.double)
            self.value_net = self.value_net.type(torch.double)
            self.replay_buffer.buffer = np.asarray(
                self.replay_buffer.buffer).astype(np.float64).tolist()
            self.soft_q_net = self.soft_q_net.type(torch.double)

    def train(self):
        """
        Trains a continuous policy by means of SAC
        """
        # Starting data
        episode = 0
        frame_count = 0
        max_episodes = args.max_episodes
        time_horizon = 250
        batch_size = args.batch_size
        # Load parameters from previous training if available
        utils.restore_data(self.restore_path)

        begin = time.time()
        while episode < max_episodes:

            if (episode % 50 == 0):  # Hack because of PyRep set position bug
                self.env.restart = True
                self.env.reset()
                self.env.restart = False

            state = self.env.reset()
            episode_reward = 0

            for step in range(time_horizon):
                if frame_count > self.min_num_steps_before_training:

                    action = self.policy_net.get_action(
                        self.state_to_tensor(state))  # .detach()
                    next_state, reward, done, self.env_info = self.env.step(
                        action * self.action_range[1])

                else:
                    action = np.random.sample(self.action_dim)
                    next_state, reward, done, self.env_info = self.env.step(
                        action * self.action_range[1])

                self.replay_buffer.push(
                    state, action, reward, next_state, done)

                if len(self.replay_buffer) > batch_size:
                    if (episode % args.num_steps_until_train) == 0:
                        for i in range(args.num_trains_per_step):
                            q_value_loss, policy_loss, value_loss = self.soft_q_update(
                                batch_size)

                state = next_state
                episode_reward += reward
                frame_count += 1

                if done:
                    break

            print("Episode = {0} | Reward = {1:.2f} | Lenght = {2:.2f}".format(
                episode, episode_reward, step))
            episode += 1

            # Saving
            if (episode % self.save_interval == 0) and (episode > 0):
                self.__save_model(episode, frame_count)

            time_elapsed = time.time() - begin
            if (episode % 100 == 0) and (episode > 0):
                print('Time elapsed so far = {0:.2f} seconds.'.format(
                    time_elapsed))

            # Logging
            if len(self.replay_buffer) > batch_size:
                self.__write_csv(episode,
                                 time_elapsed,
                                 frame_count,
                                 len(self.replay_buffer),
                                 episode_reward,
                                 value_loss,
                                 q_value_loss,
                                 policy_loss,
                                 step)

def main(args):

    env = DroneEnv(random=args.env_reset_mode, headless=True, seed=args.seed,
                   reward_function_name=args.reward_function, state=args.state)

    use_cuda = torch.cuda.is_available()

    if use_cuda and (args.use_cuda):
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # Set save/restore paths
    save_path = os.path.join('./checkpoint/', args.save_path) + '/'

    restore_path = args.restore_path or save_path
    report_folder = save_path  # save_path.split('/')[0] + '/'

    # Check if they exist
    utils.check_dir(save_path)
    if restore_path:
        utils.check_dir(restore_path)
    utils.check_dir(report_folder)

    # Preparing log csv
    if not os.path.isfile(os.path.join(report_folder, 'progress.csv')):
        print('There is no csv there')
        with open(os.path.join(report_folder, 'progress.csv'), 'w') as outcsv:
            writer = csv.writer(outcsv, delimiter=';',
                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
            writer.writerow(["Episode", "Total time (s)", "Frame",
                             "Buffer_size",
                             "Mean_Reward",
                             "value_loss", "q_value_loss", "policy_loss",
                             "episode_lenght"])

        # Network and env parameters
    action_dim = env.action_space.shape[0]
    try:
        state_dim = env.observation_space.shape[0]
    except BaseException:
        state_dim = env.observation_space
    hidden_dim = args.net_size_value
    action_range = [env.agent.action_space.low.min(
    ), env.agent.action_space.high.max()]

    sac = SAC(
        env=env,
        replay_buffer_size=args.replay_buffer_size,
        hidden_dim=hidden_dim,
        restore_path=restore_path,
        device=device,
        save_path=save_path,
        learning_rate=args.learning_rate,
        max_episodes=args.max_episodes,
        use_double=args.use_double,
        save_interval=args.save_interval)

    sac.train()

if __name__ == "__main__":
    args = argparser()

    # Setting seed
    torch.manual_seed(args.seed)
    # random.seed(a = seed)
    np.random.seed(seed=args.seed)

    main(args)


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

"""
SELFIES: a robust representation of semantically constrained graphs with an
    example application in chemistry (https://arxiv.org/abs/1905.13741)
    by Mario Krenn, Florian Haese, AkshatKuman Nigam, Pascal Friederich,
    Alan Aspuru-Guzik.

    Variational Autoencoder (VAE) for chemistry
        comparing SMILES and SELFIES representation using reconstruction
        quality, diversity and latent space validity as metrics of
        interest

information:
    ML framework: pytorch
    chemistry framework: RDKit

    get_selfie_and_smiles_encodings_for_dataset
        generate complete encoding (inclusive alphabet) for SMILES and
        SELFIES given a data file

    VAEEncoder
        fully connected, 3 layer neural network - encodes a one-hot
        representation of molecule (in SMILES or SELFIES representation)
        to latent space

    VAEDecoder
        decodes point in latent space using an RNN

    latent_space_quality
        samples points from latent space, decodes them into molecules,
        calculates chemical validity (using RDKit's MolFromSmiles), calculates
        diversity
"""

import os
import sys
import time

import numpy as np
import pandas as pd
import torch
import yaml
from rdkit import rdBase
from rdkit.Chem import MolFromSmiles
from torch import nn

import selfies as sf
from examples.vae_example.data_loader import \
    multiple_selfies_to_hot, multiple_smile_to_hot

rdBase.DisableLog('rdApp.error')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def _make_dir(directory):
    os.makedirs(directory)

def save_models(encoder, decoder, epoch):
    out_dir = './saved_models/{}'.format(epoch)
    _make_dir(out_dir)
    torch.save(encoder, '{}/E'.format(out_dir))
    torch.save(decoder, '{}/D'.format(out_dir))

class VAEEncoder(nn.Module):

    def __init__(self, in_dimension, layer_1d, layer_2d, layer_3d,
                 latent_dimension):
        """
        Fully Connected layers to encode molecule to latent space
        """
        super(VAEEncoder, self).__init__()
        self.latent_dimension = latent_dimension

        # Reduce dimension up to second last layer of Encoder
        self.encode_nn = nn.Sequential(
            nn.Linear(in_dimension, layer_1d),
            nn.ReLU(),
            nn.Linear(layer_1d, layer_2d),
            nn.ReLU(),
            nn.Linear(layer_2d, layer_3d),
            nn.ReLU()
        )

        # Latent space mean
        self.encode_mu = nn.Linear(layer_3d, latent_dimension)

        # Latent space variance
        self.encode_log_var = nn.Linear(layer_3d, latent_dimension)

    @staticmethod
    def reparameterize(mu, log_var):
        """
        This trick is explained well here:
            https://stats.stackexchange.com/a/16338
        """
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return eps.mul(std).add_(mu)

    def forward(self, x):
        """
        Pass throught the Encoder
        """
        # Get results of encoder network
        h1 = self.encode_nn(x)

        # latent space
        mu = self.encode_mu(h1)
        log_var = self.encode_log_var(h1)

        # Reparameterize
        z = self.reparameterize(mu, log_var)
        return z, mu, log_var

class VAEDecoder(nn.Module):

    def __init__(self, latent_dimension, gru_stack_size, gru_neurons_num,
                 out_dimension):
        """
        Through Decoder
        """
        super(VAEDecoder, self).__init__()
        self.latent_dimension = latent_dimension
        self.gru_stack_size = gru_stack_size
        self.gru_neurons_num = gru_neurons_num

        # Simple Decoder
        self.decode_RNN = nn.GRU(
            input_size=latent_dimension,
            hidden_size=gru_neurons_num,
            num_layers=gru_stack_size,
            batch_first=False)

        self.decode_FC = nn.Sequential(
            nn.Linear(gru_neurons_num, out_dimension),
        )

    def init_hidden(self, batch_size=1):
        weight = next(self.parameters())
        return weight.new_zeros(self.gru_stack_size, batch_size,
                                self.gru_neurons_num)

    def forward(self, z, hidden):
        """
        A forward pass throught the entire model.
        """

        # Decode
        l1, hidden = self.decode_RNN(z, hidden)
        decoded = self.decode_FC(l1)  # fully connected layer

        return decoded, hidden

def is_correct_smiles(smiles):
    """
    Using RDKit to calculate whether molecule is syntactically and
    semantically valid.
    """
    if smiles == "":
        return False

    try:
        return MolFromSmiles(smiles, sanitize=True) is not None
    except Exception:
        return False

def sample_latent_space(vae_encoder, vae_decoder, sample_len):
    vae_encoder.eval()
    vae_decoder.eval()

    gathered_atoms = []

    fancy_latent_point = torch.randn(1, 1, vae_encoder.latent_dimension,
                                     device=device)
    hidden = vae_decoder.init_hidden()

    # runs over letters from molecules (len=size of largest molecule)
    for _ in range(sample_len):
        out_one_hot, hidden = vae_decoder(fancy_latent_point, hidden)

        out_one_hot = out_one_hot.flatten().detach()
        soft = nn.Softmax(0)
        out_one_hot = soft(out_one_hot)

        out_index = out_one_hot.argmax(0)
        gathered_atoms.append(out_index.data.cpu().tolist())

    vae_encoder.train()
    vae_decoder.train()

    return gathered_atoms

def latent_space_quality(vae_encoder, vae_decoder, type_of_encoding,
                         alphabet, sample_num, sample_len):
    total_correct = 0
    all_correct_molecules = set()
    print(f"latent_space_quality:"
          f" Take {sample_num} samples from the latent space")

    for _ in range(1, sample_num + 1):

        molecule_pre = ''
        for i in sample_latent_space(vae_encoder, vae_decoder, sample_len):
            molecule_pre += alphabet[i]
        molecule = molecule_pre.replace(' ', '')

        if type_of_encoding == 1:  # if SELFIES, decode to SMILES
            molecule = sf.decoder(molecule)

        if is_correct_smiles(molecule):
            total_correct += 1
            all_correct_molecules.add(molecule)

    return total_correct, len(all_correct_molecules)

def quality_in_valid_set(vae_encoder, vae_decoder, data_valid, batch_size):
    data_valid = data_valid[torch.randperm(data_valid.size()[0])]  # shuffle
    num_batches_valid = len(data_valid) // batch_size

    quality_list = []
    for batch_iteration in range(min(25, num_batches_valid)):

        # get batch
        start_idx = batch_iteration * batch_size
        stop_idx = (batch_iteration + 1) * batch_size
        batch = data_valid[start_idx: stop_idx]
        _, trg_len, _ = batch.size()

        inp_flat_one_hot = batch.flatten(start_dim=1)
        latent_points, mus, log_vars = vae_encoder(inp_flat_one_hot)

        latent_points = latent_points.unsqueeze(0)
        hidden = vae_decoder.init_hidden(batch_size=batch_size)
        out_one_hot = torch.zeros_like(batch, device=device)
        for seq_index in range(trg_len):
            out_one_hot_line, hidden = vae_decoder(latent_points, hidden)
            out_one_hot[:, seq_index, :] = out_one_hot_line[0]

        # assess reconstruction quality
        quality = compute_recon_quality(batch, out_one_hot)
        quality_list.append(quality)

    return np.mean(quality_list).item()

def train_model(vae_encoder, vae_decoder,
                data_train, data_valid, num_epochs, batch_size,
                lr_enc, lr_dec, KLD_alpha,
                sample_num, sample_len, alphabet, type_of_encoding):
    """
    Train the Variational Auto-Encoder
    """

    print('num_epochs: ', num_epochs)

    # initialize an instance of the model
    optimizer_encoder = torch.optim.Adam(vae_encoder.parameters(), lr=lr_enc)
    optimizer_decoder = torch.optim.Adam(vae_decoder.parameters(), lr=lr_dec)

    data_train = data_train.clone().detach().to(device)
    num_batches_train = int(len(data_train) / batch_size)

    quality_valid_list = [0, 0, 0, 0]
    for epoch in range(num_epochs):

        data_train = data_train[torch.randperm(data_train.size()[0])]

        start = time.time()
        for batch_iteration in range(num_batches_train):  # batch iterator

            # manual batch iterations
            start_idx = batch_iteration * batch_size
            stop_idx = (batch_iteration + 1) * batch_size
            batch = data_train[start_idx: stop_idx]

            # reshaping for efficient parallelization
            inp_flat_one_hot = batch.flatten(start_dim=1)
            latent_points, mus, log_vars = vae_encoder(inp_flat_one_hot)

            # initialization hidden internal state of RNN (RNN has two inputs
            # and two outputs:)
            #    input: latent space & hidden state
            #    output: one-hot encoding of one character of molecule & hidden
            #    state the hidden state acts as the internal memory
            latent_points = latent_points.unsqueeze(0)
            hidden = vae_decoder.init_hidden(batch_size=batch_size)

            # decoding from RNN N times, where N is the length of the largest
            # molecule (all molecules are padded)
            out_one_hot = torch.zeros_like(batch, device=device)
            for seq_index in range(batch.shape[1]):
                out_one_hot_line, hidden = vae_decoder(latent_points, hidden)
                out_one_hot[:, seq_index, :] = out_one_hot_line[0]

            # compute ELBO
            loss = compute_elbo(batch, out_one_hot, mus, log_vars, KLD_alpha)

            # perform back propogation
            optimizer_encoder.zero_grad()
            optimizer_decoder.zero_grad()
            loss.backward(retain_graph=True)
            nn.utils.clip_grad_norm_(vae_decoder.parameters(), 0.5)
            optimizer_encoder.step()
            optimizer_decoder.step()

            if batch_iteration % 30 == 0:
                end = time.time()

                # assess reconstruction quality
                quality_train = compute_recon_quality(batch, out_one_hot)
                quality_valid = quality_in_valid_set(vae_encoder, vae_decoder,
                                                     data_valid, batch_size)

                report = 'Epoch: %d,  Batch: %d / %d,\t(loss: %.4f\t| ' \
                         'quality: %.4f | quality_valid: %.4f)\t' \
                         'ELAPSED TIME: %.5f' \
                         % (epoch, batch_iteration, num_batches_train,
                            loss.item(), quality_train, quality_valid,
                            end - start)
                print(report)
                start = time.time()

        quality_valid = quality_in_valid_set(vae_encoder, vae_decoder,
                                             data_valid, batch_size)
        quality_valid_list.append(quality_valid)

        # only measure validity of reconstruction improved
        quality_increase = len(quality_valid_list) \
                           - np.argmax(quality_valid_list)
        if quality_increase == 1 and quality_valid_list[-1] > 50.:
            corr, unique = latent_space_quality(vae_encoder, vae_decoder,
                                                type_of_encoding, alphabet,
                                                sample_num, sample_len)
        else:
            corr, unique = -1., -1.

        report = 'Validity: %.5f %% | Diversity: %.5f %% | ' \
                 'Reconstruction: %.5f %%' \
                 % (corr * 100. / sample_num, unique * 100. / sample_num,
                    quality_valid)
        print(report)

        with open('results.dat', 'a') as content:
            content.write(report + '\n')

        if quality_valid_list[-1] < 70. and epoch > 200:
            break

        if quality_increase > 20:
            print('Early stopping criteria')
            break

def compute_elbo(x, x_hat, mus, log_vars, KLD_alpha):
    inp = x_hat.reshape(-1, x_hat.shape[2])
    target = x.reshape(-1, x.shape[2]).argmax(1)

    criterion = torch.nn.CrossEntropyLoss()
    recon_loss = criterion(inp, target)
    kld = -0.5 * torch.mean(1. + log_vars - mus.pow(2) - log_vars.exp())

    return recon_loss + KLD_alpha * kld

def compute_recon_quality(x, x_hat):
    x_indices = x.reshape(-1, x.shape[2]).argmax(1)
    x_hat_indices = x_hat.reshape(-1, x_hat.shape[2]).argmax(1)

    differences = 1. - torch.abs(x_hat_indices - x_indices)
    differences = torch.clamp(differences, min=0., max=1.).double()
    quality = 100. * torch.mean(differences)
    quality = quality.detach().cpu().numpy()

    return quality

def get_selfie_and_smiles_encodings_for_dataset(file_path):
    """
    Returns encoding, alphabet and length of largest molecule in SMILES and
    SELFIES, given a file containing SMILES molecules.

    input:
        csv file with molecules. Column's name must be 'smiles'.
    output:
        - selfies encoding
        - selfies alphabet
        - longest selfies string
        - smiles encoding (equivalent to file content)
        - smiles alphabet (character based)
        - longest smiles string
    """

    df = pd.read_csv(file_path)

    smiles_list = np.asanyarray(df.smiles)

    smiles_alphabet = list(set(''.join(smiles_list)))
    smiles_alphabet.append(' ')  # for padding

    largest_smiles_len = len(max(smiles_list, key=len))

    print('--> Translating SMILES to SELFIES...')
    selfies_list = list(map(sf.encoder, smiles_list))

    all_selfies_symbols = sf.get_alphabet_from_selfies(selfies_list)
    all_selfies_symbols.add('[nop]')
    selfies_alphabet = list(all_selfies_symbols)

    largest_selfies_len = max(sf.len_selfies(s) for s in selfies_list)

    print('Finished translating SMILES to SELFIES.')

    return selfies_list, selfies_alphabet, largest_selfies_len, \
           smiles_list, smiles_alphabet, largest_smiles_len

def main():
    content = open('logfile.dat', 'w')
    content.close()
    content = open('results.dat', 'w')
    content.close()

    if os.path.exists("settings.yml"):
        settings = yaml.safe_load(open("settings.yml", "r"))
    else:
        print("Expected a file settings.yml but didn't find it.")
        return

    print('--> Acquiring data...')
    type_of_encoding = settings['data']['type_of_encoding']
    file_name_smiles = settings['data']['smiles_file']

    print('Finished acquiring data.')

    if type_of_encoding == 0:
        print('Representation: SMILES')
        _, _, _, encoding_list, encoding_alphabet, largest_molecule_len = \
            get_selfie_and_smiles_encodings_for_dataset(file_name_smiles)

        print('--> Creating one-hot encoding...')
        data = multiple_smile_to_hot(encoding_list, largest_molecule_len,
                                     encoding_alphabet)
        print('Finished creating one-hot encoding.')

    elif type_of_encoding == 1:
        print('Representation: SELFIES')
        encoding_list, encoding_alphabet, largest_molecule_len, _, _, _ = \
            get_selfie_and_smiles_encodings_for_dataset(file_name_smiles)

        print('--> Creating one-hot encoding...')
        data = multiple_selfies_to_hot(encoding_list, largest_molecule_len,
                                       encoding_alphabet)
        print('Finished creating one-hot encoding.')

    else:
        print("type_of_encoding not in {0, 1}.")
        return

    len_max_molec = data.shape[1]
    len_alphabet = data.shape[2]
    len_max_mol_one_hot = len_max_molec * len_alphabet

    print(' ')
    print(f"Alphabet has {len_alphabet} letters, "
          f"largest molecule is {len_max_molec} letters.")

    data_parameters = settings['data']
    batch_size = data_parameters['batch_size']

    encoder_parameter = settings['encoder']
    decoder_parameter = settings['decoder']
    training_parameters = settings['training']

    vae_encoder = VAEEncoder(in_dimension=len_max_mol_one_hot,
                             **encoder_parameter).to(device)
    vae_decoder = VAEDecoder(**decoder_parameter,
                             out_dimension=len(encoding_alphabet)).to(device)

    print('*' * 15, ': -->', device)

    data = torch.tensor(data, dtype=torch.float).to(device)

    train_valid_test_size = [0.5, 0.5, 0.0]
    data = data[torch.randperm(data.size()[0])]
    idx_train_val = int(len(data) * train_valid_test_size[0])
    idx_val_test = idx_train_val + int(len(data) * train_valid_test_size[1])

    data_train = data[0:idx_train_val]
    data_valid = data[idx_train_val:idx_val_test]

    print("start training")
    train_model(**training_parameters,
                vae_encoder=vae_encoder,
                vae_decoder=vae_decoder,
                batch_size=batch_size,
                data_train=data_train,
                data_valid=data_valid,
                alphabet=encoding_alphabet,
                type_of_encoding=type_of_encoding,
                sample_len=len_max_molec)

    with open('COMPLETED', 'w') as content:
        content.write('exit code: 0')

if __name__ == '__main__':
    try:
        main()
    except AttributeError:
        _, error_message, _ = sys.exc_info()
        print(error_message)


#!/usr/bin/env python
# coding=utf-8
"""
infer model
"""
import sys
import os
import numpy as np
import json
import pickle
import argparse
import time

import numpy as np
import paddle.fluid as fluid
from paddle.inference import Config
from paddle.inference import create_predictor
from paddle.inference import Predictor

from datareader import get_reader
from config import merge_configs, parse_config, print_configs

def parse_args():
    """parse_args
    """
    parser = argparse.ArgumentParser("Paddle Video infer script")
    parser.add_argument('--model_name',
                        type=str,
                        default='BaiduNet',
                        help='name of model to train.')
    parser.add_argument('--config',
                        type=str,
                        default='configs/conf.txt',
                        help='path to config file of model')
    parser.add_argument('--output', type=str, default=None, help='output path')
    parser.add_argument('--use_gpu',
                        type=bool,
                        default=True,
                        help='default use gpu.')
    parser.add_argument('--save_inference_model',
                        type=str,
                        default=None,
                        help='save inference path')
    args = parser.parse_args()
    return args

class InferModel(object):
    """lstm infer"""
    def __init__(self, cfg, name='ACTION'):
        name = name.upper()
        self.name           = name
        self.threshold      = cfg.INFER.threshold
        self.cfg            = cfg
        self.label_map      = load_class_file(cfg.MODEL.class_name_file)

    def load_inference_model(self, model_dir, use_gpu=True):
        """model_init
        """
        model_file = os.path.join(model_dir, "model")
        params_file = os.path.join(model_dir, "params")
        config = Config(model_file, params_file)
        if use_gpu:
            config.enable_use_gpu(1024)
        else:
            config.disable_gpu()
        config.switch_ir_optim(True)  # default true
        config.enable_memory_optim()
        # use zero copy
        config.switch_use_feed_fetch_ops(False)
        self.predictor = create_predictor(config)
        # build input tensor and output tensor
        self.build_input_output()

    def build_input_output(self):
        """build_input_output
        """
        input_names = self.predictor.get_input_names()
        # input
        self.input_rgb_tensor = self.predictor.get_input_handle(input_names[0])
        self.input_audio_tensor = self.predictor.get_input_handle(input_names[1])
        self.input_text_tensor = self.predictor.get_input_handle(input_names[2])

        # output
        output_names = self.predictor.get_output_names()
        self.output_tensor = self.predictor.get_output_handle(output_names[0])

    def preprocess_for_lod_data(self, input):
        """pre process"""
        input_arr = []
        input_lod = [0]
        start_lod = 0
        end_lod = 0
        for sub_item in input:
            end_lod = start_lod + len(sub_item)
            input_lod.append(end_lod)
            input_arr.extend(sub_item)
            start_lod = end_lod
        input_arr = np.array(input_arr)
        return input_arr, [input_lod]

    def predict(self):
        """predict"""
        infer_reader = get_reader(self.name, 'infer', self.cfg)
        probs = []
        video_ids = []
        label_map_inverse = {value: key for key, value in self.label_map.items()}
        for infer_iter, data in enumerate(infer_reader()):
            # video_id = [[items[-2], items[-1]] for items in data]
            rgb = [items[0] for items in data]
            audio = [items[1] for items in data]
            text = np.array([items[2] for items in data])
            videos = np.array([items[3] for items in data])

            rgb_arr, rgb_lod = self.preprocess_for_lod_data(rgb)
            audio_arr, audio_lod = self.preprocess_for_lod_data(audio)

            self.input_rgb_tensor.copy_from_cpu(rgb_arr.astype('float32'))
            self.input_rgb_tensor.set_lod(rgb_lod)

            self.input_audio_tensor.copy_from_cpu(audio_arr.astype('float32'))
            self.input_audio_tensor.set_lod(audio_lod)

            self.input_text_tensor.copy_from_cpu(text.astype('int64'))

            self.predictor.run()
            output = self.output_tensor.copy_to_cpu()
            probs.extend(list(output))
            video_ids.extend(videos)
        assert len(video_ids) == len(probs)
        result = []
        for video_id, prob in zip(video_ids, probs):
            label_idx = list(np.where(prob >= self.threshold)[0])
            result.append({
                "video_id": video_id,
                "labels": [
                    (label_map_inverse[str(idx)], float(prob[idx])) for idx in label_idx
                ]
            })
        return result

def load_class_file(class_file):
    """
    load_class_file
    """
    class_lines = open(class_file, 'r', encoding='utf8').readlines()
    class_dict = {}
    for i, line in enumerate(class_lines):
        tmp = line.strip().split('\t')
        word = tmp[0]
        index = str(i)
        if len(tmp) == 2:
            index = tmp[1]
        class_dict[word] = index
    return class_dict

def infer(args):
    """
    infer main
    """
    config = parse_config(args.config)
    infer_config = merge_configs(config, 'infer', vars(args))
    print_configs(infer_config, 'infer')
    infer_obj = InferModel(infer_config, name=args.model_name)
    infer_obj.load_inference_model(args.save_inference_model, use_gpu=args.use_gpu)
    rt = infer_obj.predict()
    if args.output:
        with open(args.output, 'w') as f:
            json.dump(rt, f, ensure_ascii=False, indent=4)

if __name__ == "__main__":
    args = parse_args()
    infer(args)


# -*- coding: utf-8 -*-
# @Author  : kevin_w

import math
import numpy as np
import torch
from modify.util import collect_action, GetRawObservations

ball_pos = [0, 0, 0]
player_pos = [0, 0, 0]
last_team = 0
last_dist_to_ball = 0
last_dist_to_goal = 0
e = np.finfo(np.float32).eps.item()

def action_modify(obs, action):
    get_obs = GetRawObservations(obs)
    team, pos, ball_direction, player = get_obs.get_ball_info()
    left_team, right_team = get_obs.get_team_position()
    active_player = left_team[get_obs.get_player()]

    # modify it if the agent repeat one single action to prevent local minimum
    # dribble the ball for a long time
    # pass the ball to each other for a long time
    # take a shoot in a long distance to the goal

    a, is_in_local_min = collect_action(action.item())
    if (team == 0) and (0 <= a <= 8) and \
            is_in_local_min:
        modified_action = torch.IntTensor([[9]])
        return modified_action

    if (team == 0) and (9 <= a <= 11) and \
            is_in_local_min:
        modified_action = torch.IntTensor([[5]])
        print('modified action2')
        return modified_action

    if (active_player[0] <= 0.4) and a == 12:
        modified_action = torch.IntTensor([[5]])
        print('modified shoot1')
        return modified_action

    if (team == 0) and (active_player[0] > 0.6) and \
            (abs(active_player[1]) < 0.25) and \
            (ball_direction[0] > 0):
        modified_action = torch.IntTensor([[12]])
        print('modified shoot2')
        return modified_action
    else:
        return action

# design a reward function based on raw info
def reward_func(obs, score, action):
    global ball_pos
    global last_team
    global last_dist_to_ball
    global last_dist_to_goal
    global player_pos
    global e
    get_obs = GetRawObservations(obs)
    team, pos, ball_direction, player = get_obs.get_ball_info()
    left_team, right_team = get_obs.get_team_position()
    team_direction = get_obs.get_team_direction()
    active_player_direction = team_direction[get_obs.get_player()]
    active_player = left_team[get_obs.get_player()]
    reward = 0
    distance = 0
    first_offense_player = max([player[0] for player in left_team])
    ball_player_distance = math.sqrt((active_player[0] - pos[0])**2 + (active_player[1] - pos[1])**2)

    # if opponents control, reward -
    if team == 1:
        reward -= 0.1
    if team == 1 and last_team != 1:
        reward -= 0.1
        last_team = 1
    if team == 0 and last_team != 0:
        # reward += 0.5
        last_team = 0

    # if ball outside the playground
    if team == 0 and \
            ((pos[0] <= -1.02) or (pos[0] >= 1.02) or (pos[1] >= 0.42) or (pos[1] <= -0.42)):
        reward -= 0.5
        print('outside punishment')

    # run to the ball and get control
    distance_to_ball = math.sqrt((pos[0] - active_player[0])**2 + (pos[1] - active_player[1])**2)
    if (last_team != 0) and (distance-last_dist_to_ball > 0.01):
        reward -= 0.1

    # action limit
    if (team != 0) and (9 <= action <= 11):
        reward -= 0.1
        print('uncontrolled punishment')
    if (team == 0) and (active_player[0] < 0.7) and \
            active_player[0] >= first_offense_player:
        if (ball_direction[0] < 0) or \
                (active_player_direction[0] < 0) or \
                ((pos[0] - ball_pos[0]) < 0.0):
            reward -= 0.1
            print('pass behind punishment')
    if (team == 0) and (active_player_direction[0] > 0.01) and \
            (ball_direction[0] > 0.01) and \
            ((action == 13) or (action == 17)):
        reward += 0.5
        print('dribble reward')
    if (active_player[0] < 0.6) and (action == 12):
        reward -= 0.1
        print('shoot punishment')
    if (team == 0) and active_player[0] > 0.6 and \
            (abs(active_player[1]) < 0.3) and \
            (action == 12):
        reward += 1
        print('shoot opportunity')

    # offense
    distance = math.sqrt((active_player[0] - 1) ** 2 + (active_player[1] - 0) ** 2)
    if pos[1] > 0:
        if (team == 0) and ((pos[0] - ball_pos[0]) > 0) and \
                ((ball_pos[1] - pos[1]) > 0) and (active_player[0] > -0.7):
            if last_dist_to_goal - distance > 0.01:
                reward += (2 - distance)
                print('move reward')
            if (active_player_direction[0] > 0.001) and \
                    ball_player_distance < 0.05:
                reward += (2 - distance)
                print('controlled reward')
    elif pos[1] < 0:
        if (team == 0) and ((pos[0] - ball_pos[0]) > 0) and \
                ((ball_pos[1] - pos[1]) < 0) and (active_player[0] > -0.7):
            if last_dist_to_goal - distance > 0.01:
                reward += (2 - distance)
                print('move reward')
            if (active_player_direction[0] > 0.001) and \
                    ball_player_distance < 0.05:
                reward += (2 - distance)
                print('controlled reward')

    # score the goal +-5
    reward += score*50

    # update record
    ball_pos = pos
    last_dist_to_ball = distance_to_ball
    last_dist_to_goal = distance
    player_pos = active_player
    return reward


# Copyright (c) 2016, NVIDIA CORPORATION.  All rights reserved.

import csv
import os

import numpy as np
import PIL.Image

class ObjectType:

    Dontcare, Car, Van, Truck, Bus, Pickup, VehicleWithTrailer, SpecialVehicle,\
        Person, Person_fa, Person_unsure, People, Cyclist, Tram, Person_Sitting,\
        Misc = range(16)

    def __init__(self):
        pass

class Bbox:

    def __init__(self, x_left=0, y_top=0, x_right=0, y_bottom=0):
        self.xl = x_left
        self.yt = y_top
        self.xr = x_right
        self.yb = y_bottom

    def area(self):
        return (self.xr - self.xl) * (self.yb - self.yt)

    def width(self):
        return self.xr - self.xl

    def height(self):
        return self.yb - self.yt

    def get_array(self):
        return [self.xl, self.yt, self.xr, self.yb]

class GroundTruthObj:

    """ This class is the data ground-truth

        #Values    Name      Description
        ----------------------------------------------------------------------------
        1    type         Class ID
        1    truncated    Float from 0 (non-truncated) to 1 (truncated), where
                          truncated refers to the object leaving image boundaries.
                          -1 corresponds to a don't care region.
        1    occluded     Integer (-1,0,1,2) indicating occlusion state:
                          -1 = unkown, 0 = fully visible,
                          1 = partly occluded, 2 = largely occluded
        1    alpha        Observation angle of object, ranging [-pi..pi]
        4    bbox         2D bounding box of object in the image (0-based index):
                          contains left, top, right, bottom pixel coordinates
        3    dimensions   3D object dimensions: height, width, length (in meters)
        3    location     3D object location x,y,z in camera coordinates (in meters)
        1    rotation_y   Rotation ry around Y-axis in camera coordinates [-pi..pi]
        1    score        Only for results: Float, indicating confidence in
                          detection, needed for p/r curves, higher is better.

        Here, 'DontCare' labels denote regions in which objects have not been labeled,
        for example because they have been too far away from the laser scanner.
    """

    # default class mappings
    OBJECT_TYPES = {
        'bus': ObjectType.Bus,
        'car': ObjectType.Car,
        'cyclist': ObjectType.Cyclist,
        'pedestrian': ObjectType.Person,
        'people': ObjectType.People,
        'person': ObjectType.Person,
        'person_sitting': ObjectType.Person_Sitting,
        'person-fa': ObjectType.Person_fa,
        'person?': ObjectType.Person_unsure,
        'pickup': ObjectType.Pickup,
        'misc': ObjectType.Misc,
        'special-vehicle': ObjectType.SpecialVehicle,
        'tram': ObjectType.Tram,
        'truck': ObjectType.Truck,
        'van': ObjectType.Van,
        'vehicle-with-trailer': ObjectType.VehicleWithTrailer}

    def __init__(self):
        self.stype = ''
        self.truncated = 0
        self.occlusion = 0
        self.angle = 0
        self.height = 0
        self.width = 0
        self.length = 0
        self.locx = 0
        self.locy = 0
        self.locz = 0
        self.roty = 0
        self.bbox = Bbox()
        self.object = ObjectType.Dontcare

    @classmethod
    def lmdb_format_length(cls):
        """
        width of an LMDB datafield returned by the gt_to_lmdb_format function.
        :return:
        """
        return 16

    def gt_to_lmdb_format(self):
        """
        For storage of a bbox ground truth object into a float32 LMDB.
         Sort-by attribute is always the last value in the array.
        """
        result = [
            # bbox in x,y,w,h format:
            self.bbox.xl,
            self.bbox.yt,
            self.bbox.xr - self.bbox.xl,
            self.bbox.yb - self.bbox.yt,
            # alpha angle:
            self.angle,
            # class number:
            self.object,
            0,
            # Y axis rotation:
            self.roty,
            # bounding box attributes:
            self.truncated,
            self.occlusion,
            # object dimensions:
            self.length,
            self.width,
            self.height,
            self.locx,
            self.locy,
            # depth (sort-by attribute):
            self.locz,
        ]
        assert(len(result) is self.lmdb_format_length())
        return result

    def set_type(self):
        self.object = self.OBJECT_TYPES.get(self.stype, ObjectType.Dontcare)

class GroundTruth:
    """
    this class loads the ground truth
    """

    def __init__(self,
                 label_dir,
                 label_ext='.txt',
                 label_delimiter=' ',
                 min_box_size=None,
                 class_mappings=None):
        self.label_dir = label_dir
        self.label_ext = label_ext  # extension of label files
        self.label_delimiter = label_delimiter  # space is used as delimiter in label files
        self._objects_all = dict()  # positive bboxes across images
        self.min_box_size = min_box_size

        if class_mappings is not None:
            GroundTruthObj.OBJECT_TYPES = class_mappings

    def update_objects_all(self, _key, _bboxes):
        if _bboxes:
            self._objects_all[_key] = _bboxes
        else:
            self._objects_all[_key] = []

    def load_gt_obj(self):
        """ load bbox ground truth from files either via the provided label directory or list of label files"""
        files = os.listdir(self.label_dir)
        files = filter(lambda x: x.endswith(self.label_ext), files)
        if len(files) == 0:
            raise RuntimeError('error: no label files found in %s' % self.label_dir)
        for label_file in files:
            objects_per_image = list()
            with open(os.path.join(self.label_dir, label_file), 'rb') as flabel:
                for row in csv.reader(flabel, delimiter=self.label_delimiter):
                    if len(row) == 0:
                        # This can happen when you open an empty file
                        continue
                    if len(row) < 15:
                        raise ValueError('Invalid label format in "%s"'
                                         % os.path.join(self.label_dir, label_file))

                    # load data
                    gt = GroundTruthObj()
                    gt.stype = row[0].lower()
                    gt.truncated = float(row[1])
                    gt.occlusion = int(row[2])
                    gt.angle = float(row[3])
                    gt.bbox.xl = float(row[4])
                    gt.bbox.yt = float(row[5])
                    gt.bbox.xr = float(row[6])
                    gt.bbox.yb = float(row[7])
                    gt.height = float(row[8])
                    gt.width = float(row[9])
                    gt.length = float(row[10])
                    gt.locx = float(row[11])
                    gt.locy = float(row[12])
                    gt.locz = float(row[13])
                    gt.roty = float(row[14])
                    gt.set_type()
                    box_dimensions = [gt.bbox.xr - gt.bbox.xl, gt.bbox.yb - gt.bbox.yt]
                    if self.min_box_size is not None:
                        if not all(x >= self.min_box_size for x in box_dimensions):
                            # object is smaller than threshold => set to "DontCare"
                            gt.stype = ''
                            gt.object = ObjectType.Dontcare
                    objects_per_image.append(gt)
                key = os.path.splitext(label_file)[0]
                self.update_objects_all(key, objects_per_image)

    @property
    def objects_all(self):
        return self._objects_all

# return the # of pixels remaining in a

def pad_bbox(arr, max_bboxes=64, bbox_width=16):
    if arr.shape[0] > max_bboxes:
        raise ValueError(
            'Too many bounding boxes (%d > %d)' % arr.shape[0], max_bboxes
        )
    # fill remainder with zeroes:
    data = np.zeros((max_bboxes + 1, bbox_width), dtype='float')
    # number of bounding boxes:
    data[0][0] = arr.shape[0]
    # width of a bounding box:
    data[0][1] = bbox_width
    # bounding box data. Merge nothing if no bounding boxes exist.
    if arr.shape[0] > 0:
        data[1:1 + arr.shape[0]] = arr

    return data

def bbox_to_array(arr, label=0, max_bboxes=64, bbox_width=16):
    """
    Converts a 1-dimensional bbox array to an image-like
    3-dimensional array CHW array
    """
    arr = pad_bbox(arr, max_bboxes, bbox_width)
    return arr[np.newaxis, :, :]

def bbox_overlap(abox, bbox):
    # the abox box
    x11 = abox[0]
    y11 = abox[1]
    x12 = abox[0] + abox[2] - 1
    y12 = abox[1] + abox[3] - 1

    # the closer box
    x21 = bbox[0]
    y21 = bbox[1]
    x22 = bbox[0] + bbox[2] - 1
    y22 = bbox[1] + bbox[3] - 1

    overlap_box_x2 = min(x12, x22)
    overlap_box_x1 = max(x11, x21)
    overlap_box_y2 = min(y12, y22)
    overlap_box_y1 = max(y11, y21)

    # make sure we preserve any non-bbox components
    overlap_box = list(bbox)
    overlap_box[0] = overlap_box_x1
    overlap_box[1] = overlap_box_y1
    overlap_box[2] = overlap_box_x2 - overlap_box_x1 + 1
    overlap_box[3] = overlap_box_y2 - overlap_box_y1 + 1

    xoverlap = max(0, overlap_box_x2 - overlap_box_x1)
    yoverlap = max(0, overlap_box_y2 - overlap_box_y1)
    overlap_pix = xoverlap * yoverlap

    return overlap_pix, overlap_box

def pad_image(img, padding_image_height, padding_image_width):
    """
    pad a single image to the specified dimensions
    """
    src_width = img.size[0]
    src_height = img.size[1]

    if padding_image_width < src_width:
        raise ValueError("Source image width %d is greater than padding width %d" % (src_width, padding_image_width))

    if padding_image_height < src_height:
        raise ValueError("Source image height %d is greater than padding height %d" %
                         (src_height, padding_image_height))

    padded_img = PIL.Image.new(
        img.mode,
        (padding_image_width, padding_image_height),
        "black")
    padded_img.paste(img, (0, 0))  # copy to top-left corner

    return padded_img

def resize_bbox_list(bboxlist, rescale_x=1, rescale_y=1):
        # this is expecting x1,y1,w,h:
    bboxListNew = []
    for bbox in bboxlist:
        abox = bbox
        abox[0] *= rescale_x
        abox[1] *= rescale_y
        abox[2] *= rescale_x
        abox[3] *= rescale_y
        bboxListNew.append(abox)
    return bboxListNew


from __future__ import absolute_import, division
import numpy as np
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
# import fitlog
# fitlog.commit(__file__)

# __all__ = ['BalancedLoss', 'FocalLoss', 'GHMCLoss', 'OHNMLoss','SmoothL1']
# For rpn_right now , only
__all__ = ['BalancedCELoss', 'SmoothL1']

def log_sigmoid(x):
    # for x > 0: 0 - log(1 + exp(-x))
    # for x < 0: x - log(1 + exp(x))
    # for x = 0: 0 (extra term for gradient stability)
    return torch.clamp(x, max=0) - torch.log(1 + torch.exp(-torch.abs(x))) + \
        0.5 * torch.clamp(x, min=0, max=0)

def log_minus_sigmoid(x):
    # for x > 0: -x - log(1 + exp(-x))
    # for x < 0:  0 - log(1 + exp(x))
    # for x = 0: 0 (extra term for gradient stability)
    return torch.clamp(-x, max=0) - torch.log(1 + torch.exp(-torch.abs(x))) + \
        0.5 * torch.clamp(x, min=0, max=0)

class BalancedCELoss(nn.Module):

    def __init__(self):
        super(BalancedCELoss, self).__init__()
        # self.neg_weight = neg_weight

    def forward(self, inputs, target , num_pos , num_neg):  # inputs [N , 1805 , H ] N = 8 , H = 2  target [N , 1805]
        loss_all = []
        for batch_id in range(target.shape[0]):
            min_pos = min(len(np.where(target[batch_id].cpu() == 1)[0]), num_pos)
            min_neg = int(min(len(np.where(target[batch_id].cpu() == 1)[0]) * num_neg / num_pos, num_neg))
            pos_index = np.where(target[batch_id].cpu() == 1)[0].tolist()
            neg_index = np.where(target[batch_id].cpu() == 0)[0].tolist()

            pos_index_random = random.sample(pos_index, min_pos)
            if len(pos_index) > 0:
                pos_loss_id_final = F.cross_entropy(input=inputs[batch_id][pos_index_random],
                                                     target=target[batch_id][pos_index_random], reduction='none')
            else:
                pos_loss_id_final = torch.FloatTensor([0]).cuda()

            neg_index_random = random.sample(neg_index, min_neg)
            if len(neg_index) > 0:
                neg_loss_id_final = F.cross_entropy(input=inputs[batch_id][neg_index_random],
                                                     target=target[batch_id][neg_index_random], reduction='none')
            else:
                neg_loss_id_final = torch.FloatTensor([0]).cuda()

            loss_id = (pos_loss_id_final.mean() + neg_loss_id_final.mean()) / 2
            loss_all.append(loss_id)
        # print(loss_all)
        # print(loss_all)
        return torch.stack(loss_all).mean()

class SmoothL1(nn.Module):
    def __init__(self):
        super(SmoothL1, self).__init__()
    def forward(self, inputs, target , label , num_pos):  # inputs [N , 1805 , H ] N = 8 , H = 4  target [N , 1805 , 4]  label : [pos , neg , ]
        loss_all = []
        for batch_id in range(target.shape[0]):
            pos_index = np.where(label[batch_id].cpu() == 1)[0]   # 仅对证样本进行计算
            min_pos = min(len(pos_index), num_pos)
            pos_index = random.sample(pos_index.tolist(), min_pos)
            if len(pos_index) > 0:
                loss_id = F.smooth_l1_loss(inputs[batch_id][pos_index], target[batch_id][pos_index])
                # print(loss_id)
            else:
                loss_id = torch.FloatTensor([0]).cuda()[0]
            loss_all.append(loss_id.mean())
        # print(loss_all)
        return torch.stack(loss_all).mean()

class FocalLoss(nn.Module):

    def __init__(self, gamma=2):
        super(FocalLoss, self).__init__()
        self.gamma = gamma

    def forward(self, input, target):
        pos_log_sig = log_sigmoid(input)
        neg_log_sig = log_minus_sigmoid(input)

        prob = torch.sigmoid(input)
        pos_weight = torch.pow(1 - prob, self.gamma)
        neg_weight = torch.pow(prob, self.gamma)

        loss = -(target * pos_weight * pos_log_sig + \
            (1 - target) * neg_weight * neg_log_sig)

        avg_weight = target * pos_weight + (1 - target) * neg_weight
        loss /= avg_weight.mean()

        return loss.mean()

class GHMCLoss(nn.Module):

    def __init__(self, bins=30, momentum=0.5):
        super(GHMCLoss, self).__init__()
        self.bins = bins
        self.momentum = momentum
        self.edges = [t / bins for t in range(bins + 1)]
        self.edges[-1] += 1e-6
        if momentum > 0:
            self.acc_sum = [0.0 for _ in range(bins)]

    def forward(self, input, target):
        edges = self.edges
        mmt = self.momentum
        weights = torch.zeros_like(input)

        # gradient length
        g = torch.abs(input.sigmoid().detach() - target)

        tot = input.numel()
        n = 0
        for i in range(self.bins):
            inds = (g >= edges[i]) & (g < edges[i + 1])
            num_in_bin = inds.sum().item()
            if num_in_bin > 0:
                if mmt > 0:
                    self.acc_sum[i] = mmt * self.acc_sum[i] \
                        + (1 - mmt) * num_in_bin
                    weights[inds] = tot / self.acc_sum[i]
                else:
                    weights[inds] = tot / num_in_bin
                n += 1
        if n > 0:
            weights /= weights.mean()

        loss = F.binary_cross_entropy_with_logits(
            input, target, weights, reduction='sum') / tot

        return loss

class OHNMLoss(nn.Module):

    def __init__(self, neg_ratio=3.0):
        super(OHNMLoss, self).__init__()
        self.neg_ratio = neg_ratio

    def forward(self, input, target):
        pos_logits = input[target > 0]
        pos_labels = target[target > 0]

        neg_logits = input[target == 0]
        neg_labels = target[target == 0]

        pos_num = pos_logits.numel()
        neg_num = int(pos_num * self.neg_ratio)
        neg_logits, neg_indices = neg_logits.topk(neg_num)
        neg_labels = neg_labels[neg_indices]

        loss = F.binary_cross_entropy_with_logits(
            torch.cat([pos_logits, neg_logits]),
            torch.cat([pos_labels, neg_labels]),
            reduction='mean')

        return loss


import matplotlib
matplotlib.use('Agg')
import os, sys
import yaml
from argparse import ArgumentParser
from tqdm import tqdm

import imageio
from imageio import mimread
import numpy as np
from skimage.transform import resize
from skimage import img_as_ubyte
from skimage import io
import torch
import torch.nn.functional as F

from modules.generator import OcclusionAwareGenerator
from modules.keypoint_detector import KPDetector
from modules.tdmm_estimator import TDMMEstimator
from modules.flame_config import cfg as flame_cfg

from animate import normalize_kp
from scipy.spatial import ConvexHull

import time
import cv2
import face_alignment
import glob

def bbox_increase(bboxes, img_w, img_h, ratio=1.5):
    new_bboxes = bboxes.copy().astype(float)

    w = (new_bboxes[2] - new_bboxes[0])
    h = (new_bboxes[3] - new_bboxes[1])

    new_bboxes[1] = np.where(w > h, new_bboxes[1] - (w - h) / 2, new_bboxes[1])
    new_bboxes[3] = np.where(w > h, new_bboxes[3] + (w - h) / 2, new_bboxes[3])

    new_bboxes[0] = np.where(w < h, new_bboxes[0] - (h - w) / 2, new_bboxes[0])
    new_bboxes[2] = np.where(w < h, new_bboxes[2] + (h - w) / 2, new_bboxes[2])

    center = np.stack(((new_bboxes[0] + new_bboxes[2])/2, (new_bboxes[1] + new_bboxes[3])/2), axis=0)
    img_size = np.where(w > h, w, h)
    img_size = img_size * ratio

    new_bboxes[0] = np.maximum((center[0] - img_size / 2.0), 0)
    new_bboxes[1] = np.maximum((center[1] - img_size / 2.0), 0)
    new_bboxes[2] = np.minimum(img_w, (center[0] + img_size / 2.0))
    new_bboxes[3] = np.minimum(img_h, (center[1] + img_size / 2.0))

    return new_bboxes

if sys.version_info[0] < 3:
    raise Exception("You must use Python 3 or higher. Recommended version is Python 3.7")

def partial_state_dict_load(module, state_dict):
    own_state = module.state_dict()
    for name, param in state_dict.items():
        if name not in own_state:
            continue
        if isinstance(param, torch.nn.Parameter):
            # backwards compatibility for serialized parameters
            param = param.data
        own_state[name].copy_(param)

def load_checkpoints(blend_scale, config_path, checkpoint_path, cpu=False):

    with open(config_path) as f:
        config = yaml.load(f)

    # init generator
    generator = OcclusionAwareGenerator(blend_scale=blend_scale,
                                        **config['model_params']['generator_params'],
                                        **config['model_params']['common_params'])
    # init kp_detector
    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
                             **config['model_params']['common_params'])
    # init tdmm estimator
    tdmm = TDMMEstimator(flame_cfg)

    checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
    partial_state_dict_load(generator, checkpoint['generator'])
    kp_detector.load_state_dict(checkpoint['kp_detector'])
    tdmm.load_state_dict(checkpoint['tdmm'])

    if not cpu:
        generator = generator.cuda()
        kp_detector = kp_detector.cuda()
        tdmm = tdmm.cuda()

    generator.eval()
    kp_detector.eval()
    tdmm.eval()

    return generator, kp_detector, tdmm

def load_face_parser(cpu=False):
    from face_parsing.model import BiSeNet

    face_parser = BiSeNet(n_classes=19)
    face_parser.load_state_dict(torch.load('face_parsing/cp/79999_iter.pth', map_location=torch.device('cpu')), strict=False)

    if not cpu:
       face_parser.cuda()

    face_parser.eval()

    mean = torch.Tensor(np.array([0.485, 0.456, 0.406], dtype=np.float32)).view(1, 3, 1, 1)
    std = torch.Tensor(np.array([0.229, 0.224, 0.225], dtype=np.float32)).view(1, 3, 1, 1)

    if not cpu:
        face_parser.mean = mean.cuda()
        face_parser.std = std.cuda()
    else:
        face_parser.mean = mean
        face_parser.std = std

    return face_parser

def faceswap(opt, fa, generator, kp_detector, tdmm):
    source_image = np.array(io.imread(opt.source_image_pth))
    if source_image.shape[2] == 4:
        source_image = source_image[..., 0:3]
    driving_video = np.array(mimread(opt.driving_video_pth, memtest=False))

    if opt.use_detection:
        detection = fa.face_detector.detect_from_image(source_image)
        if len(detection) == 0:
            raise ValueError('No faces detected in source images')
        source_bbox = bbox_increase(detection[0][0:4], source_image.shape[1], source_image.shape[0])
        source_bbox = source_bbox.astype(np.int32)
        source = torch.tensor(source_image[np.newaxis, source_bbox[1]:source_bbox[3], source_bbox[0]:source_bbox[2]].astype(np.float32)).permute(0, 3, 1, 2)
    else:
        source = torch.tensor(source_image[np.newaxis, ...].astype(np.float32)).permute(0, 3, 1, 2)

    source = source / 255.0
    if not opt.cpu:
        source = source.cuda()
    source = F.interpolate(source, size=(256, 256), mode='bilinear', align_corners=True)
    kp_source = kp_detector(source)

    with torch.no_grad():
        predictions = []
        for i in tqdm(range(driving_video.shape[0])):
            driving_image = driving_video[i]

            if opt.use_detection:
                detection = fa.face_detector.detect_from_image(driving_image)
                if len(detection) == 0:
                    raise ValueError('No faces detected in source images')
                driving_bbox = bbox_increase(detection[0][0:4], driving_image.shape[1], driving_image.shape[0])
                driving_bbox = driving_bbox.astype(np.int32)
                driving = torch.tensor(driving_image[np.newaxis, driving_bbox[1]:driving_bbox[3], driving_bbox[0]:driving_bbox[2]].astype(np.float32)).permute(0, 3, 1, 2)
            else:
                driving = torch.tensor(driving_image[np.newaxis, ...].astype(np.float32)).permute(0, 3, 1, 2)

            driving = driving / 255.0
            if not opt.cpu:
                driving = driving.cuda()
            driving = F.interpolate(driving, size=(256, 256), mode='bilinear', align_corners=True)
            kp_driving = kp_detector(driving)

            # for face swap
            if face_parser is not None:
                blend_mask = F.interpolate(driving, size=(512, 512))
                blend_mask = (blend_mask - face_parser.mean) / face_parser.std
                blend_mask = face_parser(blend_mask)
                blend_mask = torch.softmax(blend_mask[0], dim=1)
            else:
                blend_mask = None

            blend_mask = blend_mask[:, opt.swap_index].sum(dim=1, keepdim=True)
            if opt.hard:
                blend_mask = (blend_mask > 0.5).type(blend_mask.type())

            # 3DMM rendering
            source_codedict = tdmm.encode(source)
            driving_codedict = tdmm.encode(driving)
            source_verts, source_transformed_verts, source_ldmk_2d = tdmm.decode_flame(source_codedict)
            driving_verts, driving_transformed_verts, driving_ldmk_2d = tdmm.decode_flame(driving_codedict)
            source_albedo = tdmm.extract_texture(source, source_transformed_verts, with_eye=opt.with_eye)
            render_ops = tdmm.render(source_transformed_verts, driving_transformed_verts, source_albedo)

            out = generator(source, kp_source=kp_source, kp_driving=kp_driving, render_ops=render_ops,
                            blend_mask=blend_mask, driving_image=driving, driving_features=driving_codedict)

            if opt.use_detection:
                # fit to the original video
                bbox_w, bbox_h = driving_bbox[2]-driving_bbox[0], driving_bbox[3]-driving_bbox[1]
                prediction = F.interpolate(out['prediction'], size=(bbox_h, bbox_w), mode='bilinear', align_corners=True)
                driving_image[driving_bbox[1]:driving_bbox[3], driving_bbox[0]:driving_bbox[2], :] = np.transpose(prediction.data.cpu().numpy(), [0, 2, 3, 1])[0] * 255.0
                predictions.append(driving_image)
            else:
                predictions.append((np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0] * 255.0).astype(np.uint8))

    imageio.mimsave(opt.result_video_pth, predictions, fps=25)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--config", required=True, help="path to config")
    parser.add_argument("--checkpoint", default='', help="path to checkpoint to restore")
    parser.add_argument("--with_eye", action="store_true", help="use eye part for extracting texture")
    parser.add_argument("--source_image_pth", default='', help="path to source image")
    parser.add_argument("--driving_video_pth", default='', help="path to driving video")
    parser.add_argument("--result_video_pth", default='result.mp4', help="path to output")

    parser.add_argument("--swap_index", default="1,2,3,4,5,6,7,8,9,10,11,12,13,14,15", type=lambda x: list(map(int, x.split(','))),
                        help='index of swaped parts')
    parser.add_argument("--hard", action="store_true", help="use hard segmentation labels for blending")

    parser.add_argument("--cpu", dest="cpu", action="store_true", help="cpu mode.")
    parser.add_argument("--use_detection", action="store_true", help="use detected bbox")

    opt = parser.parse_args()

    blend_scale = (256 / 4) / 512
    generator, kp_detector, tdmm = load_checkpoints(blend_scale=blend_scale, config_path=opt.config, checkpoint_path=opt.checkpoint, cpu=opt.cpu)

    face_parser = load_face_parser(opt.cpu)
    print("face_parser is loaded!")

    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, face_detector='sfd',
                                      device='cpu' if opt.cpu else 'cuda')
    faceswap(opt, fa, generator, kp_detector, tdmm)


import cv2
import tflite_runtime.interpreter as tflite
import numpy as np

CAMERA_WIDTH = 640
CAMERA_HEIGHT = 480

cap = cv2.VideoCapture(1)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT)

DETECTION_MODEL_PATH = 'DetectionModel.tflite'
CLASSIFICATION_MODEL_PATH = 'ClassificationModel.tflite'

colors = [(47,255,173), (255,255,0), (71,99,255), (211,85,186), (180,105,255),
(0,215,255), (87,139,46), (255,105,65), (45,82,160), (205,250,255)]

def set_input_tensor(interpreter, image):
    tensor_index = interpreter.get_input_details()[0]['index']
    input_tensor = interpreter.tensor(tensor_index)()[0]
    input_tensor[:,:] = image

def get_output_tensor(interpreter, index):
    output_details = interpreter.get_output_details()[index]
    tensor = np.squeeze(interpreter.get_tensor(output_details['index']))
    return tensor

def detect_soms(interpreter, image, thd=0.9):
    set_input_tensor(interpreter, image)
    interpreter.invoke()

    boxes = get_output_tensor(interpreter, 0)
    classes = get_output_tensor(interpreter, 1)
    scores = get_output_tensor(interpreter, 2)
    count = int(get_output_tensor(interpreter, 3))

    results = []
    for i in range(count):
        if scores[i] >= thd:
            results.append({'box': boxes[i]})
        else:
            break

    return results

def classify_som(interpreter, image, top_k=1):
    set_input_tensor(interpreter, image)
    interpreter.invoke()
    output_details = interpreter.get_output_details()[0]
    output = np.squeeze(interpreter.get_tensor(output_details['index']))

    if output_details['dtype']==np.uint8:
        scale, zero_point = output_details['quantization']
        output = scale*(output-zero_point)

    ordered = np.argpartition(-output, top_k)

    return [(i, output[i]) for i in ordered[:top_k]]

def main():
    modelDetect = tflite.Interpreter(DETECTION_MODEL_PATH)
    modelDetect.allocate_tensors()

    modelClassify = tflite.Interpreter(CLASSIFICATION_MODEL_PATH)
    modelClassify.allocate_tensors()

    _, input_height_detect, input_width_detect,_ = modelDetect.get_input_details()[0]['shape']
    _, input_height_classify, input_width_classify,_ = modelClassify.get_input_details()[0]['shape']

    favor_txt = 'loading...'

    while 1:
        _, frame = cap.read()
        frameshow = frame.copy()

        #detection part
        imgrz = cv2.resize(frame.copy()[:,:,[2,1,0]], (input_height_detect, input_width_detect))

        bboxes = detect_soms(modelDetect, imgrz)

        if len(bboxes)==0:
            favor_txt = 'loading...'

        for i in range(len(bboxes)):
            bb = bboxes[i]['box']
            ymin, xmin, ymax, xmax = bb
            xmin = int(xmin*CAMERA_WIDTH)
            ymin = int(ymin*CAMERA_HEIGHT)
            xmax = int(xmax*CAMERA_WIDTH)
            ymax = int(ymax*CAMERA_HEIGHT)

            #image preparation part
            imgcrop = frame.copy()[ymin:ymax, xmin:xmax, [2,1,0]]
            imgcrop = cv2.resize(imgcrop, (input_height_classify, input_width_classify))

            #prediction part
            favor = classify_som(modelClassify, imgcrop)

            if favor[0][0]==0 and favor[0][1]>=0.8:
                favor_txt = 'sweet'
            elif favor[0][0]==1 and favor[0][1]>=0.8:
                favor_txt = 'sour'
#            else:
#                favor_txt = 'not sure'

            frameshow = cv2.rectangle(frameshow, (xmin, ymin), (xmax, ymax), colors[i%10], 2)
            frameshow = cv2.putText(frameshow, favor_txt, (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX,
                                    1, colors[i%10], 2)

        cv2.imshow('somsom', frameshow)

        if cv2.waitKey(1)==27:
            cv2.destroyAllWindows()
            break

if __name__=='__main__':
    main()


#!/usr/bin/python

#By Sun Jinyuan, 2021

import pandas as pd
import numpy as np

class feature:
    __a = """AAa H V P IP HT ST GSI F0 F1 F2 C B EC
    A -0.171 −0.677 −0.680 −0.170 0.900 −0.476 −0.350 −0.044 −0.234 −0.269 0.587 −0.099 0.829
    D -0.767 −0.281 −0.417 −0.900 −0.155 −0.635 −0.213 −0.103 0.900 0.014 −0.475 −0.082 0.247
    C 0.508 −0.359 −0.329 −0.114 −0.652 0.476 −0.140 −0.642 −0.773 −0.035 −0.433 0.094 −0.388
    E -0.696 −0.058 −0.241 −0.868 0.900 −0.582 −0.230 0.347 0.480 0.021 −0.900 0.105 0.565
    F 0.646 0.412 0.373 −0.272 0.155 0.318 0.363 −0.863 −0.504 −0.113 −0.673 0.721 0.035
    G -0.342 −0.900 −0.900 −0.179 −0.900 −0.900 −0.900 0.701 0.527 −0.050 0.378 −0.900 0.829
    H −0.271 0.138 0.110 0.195 −0.031 −0.106 0.384 −0.480 −0.186 −0.255 −0.297 0.115 −0.088
    I 0.652 −0.009 −0.066 −0.186 0.155 0.688 0.900 −0.332 −0.662 −0.411 −0.288 0.879 −0.900
    K −0.889 0.163 0.066 0.727 0.279 −0.265 −0.088 0.339 0.844 0.900 −0.375 0.317 0.547
    L 0.596 −0.009 −0.066 −0.186 0.714 −0.053 0.213 −0.590 −0.115 −0.064 −0.288 0.879 0.865
    M 0.337 0.087 0.066 −0.262 0.652 −0.001 0.110 −0.738 −0.900 −0.893 −0.205 0.370 0.724
    N −0.674 −0.243 −0.329 −0.075 −0.403 −0.529 −0.213 0.516 0.242 0.000 −0.166 0.031 0.265
    P 0.055 −0.294 −0.900 −0.010 −0.900 0.106 0.247 0.059 0.868 0.014 0.900 0.487 0.212
    Q −0.464 −0.020 −0.110 −0.276 0.528 −0.371 −0.230 0.870 0.416 −0.319 −0.403 0.192 0.529
    R −0.900 0.466 0.373 0.900 0.528 −0.371 0.105 −0.066 0.416 −0.206 0.430 0.175 −0.106
    S −0.364 −0.544 −0.637 −0.265 −0.466 −0.212 −0.337 0.900 0.575 −0.050 −0.024 −0.300 0.600
    T −0.199 −0.321 −0.417 −0.288 −0.403 0.212 0.402 0.192 0.599 0.028 −0.212 0.323 0.406
    V 0.331 −0.232 −0.285 −0.191 −0.031 0.900 0.677 −0.480 −0.385 −0.120 −0.127 0.896 0.794
    W 0.900 0.900 0.900 −0.209 0.279 0.529 0.479 −0.900 −0.464 −0.900 −0.074 0.900 0.900
    Y 0.188 0.541 0.417 −0.274 −0.155 0.476 0.363 −0.634 −0.361 −0.659 −0.738 0.546 0.582"""

    vloumn = {"A":88.6,"R":173.4,"N":114.1,"D":111.1,"C":108.5,
                "Q":143.8,"E":138.4,"G":60.1,"H":153.2,"I":166.7,
                "L":166.7,"K":168.6,"M":162.9,"F":189.9,"P":112.7,
                "S":89.0,"T":116.1,"W":227.8,"Y":193.6,"V":140.0}

    hydropathy_index = {"R":-2.5,"K":-1.5,"D":-0.9,"Q":-0.85,"N":-0.78,
                          "E":-0.74,"H":-0.4,"S":-0.18,"T":-0.05,"P":0.12,
                          "Y":0.26,"C":0.29,"G":0.48,"A":0.62,"M":0.64,
                          "W":0.81,"L":1.1,"V":1.1, "F":1.2,"I":1.4}

    def get_al_dd(self):

        al_dd = {"A":{},"D":{},"C":{},"E":{},"F":{},
                 "G":{},"H":{},"I":{},"K":{},"L":{},
                 "M":{},"N":{},"P":{},"Q":{},"R":{},
                 "S":{},"T":{},"V":{},"W":{},"Y":{}}

        for line in self.__a.replace("−","-").split("\n"):
            lst = line.split()
            if lst[0] == "AAa":
                name_lst = lst[1:]
            else:
                al_dd[lst[0]] = dict(zip(name_lst,lst[1:]))

        #al_df = pd.DataFrame(al_dd).T
        return al_dd

feature = feature()
al_dd = feature.get_al_dd()
vloumn = feature.vloumn
hydropathy_index = feature.hydropathy_index

alphabet = "ARNDCQEGHILKMFPSTWYV-"
states = len(alphabet)
a2n = {}
for a,n in zip(alphabet,range(states)):
    a2n[a] = n

def get_MSA(seqa3m):
    import string
    rm_lc = str.maketrans(dict.fromkeys(string.ascii_lowercase))
    ali_dict = {}
    a3mfile = open(seqa3m)
    for line in a3mfile:
        if line[0] == ">":
            head = line.strip()
            ali_dict[head] = ""
            #print(line.strip())
        else:
            line = line.translate(rm_lc).strip()
            lst = []
            for x in line:
                lst.append(x)
            ali_dict[head] = lst
            #print(lst)
    df = pd.DataFrame(ali_dict)
    seq_length = df.shape[0]
    seq_num = df.shape[1]
    df_msa = df.T

    alphabet = "ARNDCQEGHILKMFPSTWYV-"
    p_msa = np.zeros([21,seq_length])
    for p in range(seq_length):
        for num,AA in enumerate(alphabet):
            if AA != "-":
                count = df_msa[p].value_counts().get(AA)
                if count == None:
                    p_msa[num][p] = 0
                else:
                    p_msa[num][p] = count/seq_num
                    #print(num,p,count/seq_num)
        p = p + 1

    return p_msa
def getSeq(fasta):
    fasta = open(fasta)
    seq = ''
    for line in fasta:
        if line[0] == ">":
            continue
        else:
            seq += line.strip().replace("X","A")
    return seq

def getStatic(wpm):
    dP = 0.5*(-float(al_dd[wpm[0]]['P'])+float(al_dd[wpm[2]]['P']))
    dF1 = -0.5*(-float(al_dd[wpm[0]]['F1'])+float(al_dd[wpm[2]]['F1']))
    dB = 0.5*(-float(al_dd[wpm[0]]['B'])+float(al_dd[wpm[2]]['B']))
    dST = 0.5*(-float(al_dd[wpm[0]]['ST'])+float(al_dd[wpm[2]]['ST']))
    dHT = 0.5*(-float(al_dd[wpm[0]]['HT'])+float(al_dd[wpm[2]]['HT']))
    dV = -0.007*(vloumn[wpm[0]]-vloumn[wpm[2]])
    dH = -0.25*(hydropathy_index[wpm[0]]-hydropathy_index[wpm[2]])
    return [dP, dF1, dB, dST, dHT, dV, dH]
def getStatical(msa,v_out,w_out,seq,mutation):
    wild = mutation[0]
    pos = int(mutation[1])
    mut = mutation[2]
    wild_score = []
    mut_score = []
    i = 0
    wild_v = v_out[pos-1][a2n[wild]]
    mut_v = v_out[pos-1][a2n[mut]]
    for aa in seq:
        wild_score.append(w_out[i][pos-1][:-1].reshape(21,21)[a2n[aa]][a2n[wild]])
        mut_score.append(w_out[i][pos-1][:-1].reshape(21,21)[a2n[aa]][a2n[mut]])
        i += 1

    return [-0.1*(wild_v-mut_v),
            -0.5*(msa[a2n[wild]][pos-1]-msa[a2n[mut]][pos-1]),
            -0.5*(max(wild_score) - max(mut_score)),
            0.5*(min(wild_score) - min(mut_score)),
            -0.15*(sum(wild_score) - sum(mut_score)),
            -0.1*(np.linalg.norm((np.array(wild_score)), ord=1)-np.linalg.norm((np.array(mut_score)), ord=1)),
            -0.5*(np.linalg.norm((np.array(wild_score)), ord=2)-np.linalg.norm((np.array(mut_score)), ord=2))]

if __name__ == '__main__':
    feature = feature()
    al_dd = feature.get_al_dd()
    seq = feature.getSeq(fasta)
    msa = feature.get_MSA(seqa3m)
    static_features = feature.getStatic(wpm)
    statical_features = feature.getStatical(msa,v_out,w_out,seq,mutation)


## Training code for the Imitation learning model
## Use the Driving simulator attached to this code to drive a car
## Drive nicely as the Imitation model will try to Imitate you.
## After drving, close the simulator. It will save a log of your driving data
## The training code will access this data for training
## For the model to perform well, one needs to drive
# atleast for more than an hour.
## Attaching a sample training dataset, so that one can train directly
#  without having to generate the expert training data
## The training will save a model in the folder new_training_model

import os
import cv2
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import numpy as np
import pandas as pd
from torch.utils import data
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import MultiStepLR
import scipy
from scipy import signal

# ## Helper functions
def toDevice(datas, device):
    """Enable cuda."""
    imgs, angles = datas
    return imgs.float().to(device), angles.float().to(device)

def augment(img, angle):
    """Data augmentation."""
    #name = dataroot
    current_image = cv2.imread(img)
    #cropping image to remove the sky
    current_image = current_image[60::,::]
    if np.random.rand() < 0.5:
        # data augmentation
        current_image = cv2.flip(current_image, 1)
        angle = angle * -1.0

    return current_image, angle

def change_bright(img):
    # convert rgb to hsv
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    rand = np.random.uniform(0.5,1.0)
    # increase of brightness in random images
    hsv[:,:,2] = rand*hsv[:,:2]
    # convert back to rgb
    new_img = cv2.cv2.cvtColor(hsv,cv2.COLOR_HSV2RGB)
    return new_img

def load_data(img_paths,steers, test_size):
    """Load training data and train validation split"""

    # Divide the data into training set and validation set
    data_df = pd.DataFrame({'center':img_paths,'steering':steers})
    train_len = int(test_size * len(steers))
    valid_len = len(steers) - train_len
    trainset, valset = data.random_split(
        data_df.values.tolist(), lengths=[train_len, valid_len])

    return trainset, valset

## Create dataset
class TripletDataset(data.Dataset):

    def __init__(self,dataroot,samples, transform=None):
        self.samples = samples
        self.dataroot = dataroot
        self.transform = transform

    def __getitem__(self, index):
        batch_samples = self.samples[index]
        steering_angle = float(batch_samples[1])
        center_img, steering_angle_center = augment(batch_samples[0],steering_angle)
        return (center_img, steering_angle_center)

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

# ## Get data loader
def data_loader(dataroot, trainset, valset, batch_size, shuffle, num_workers):
    """dataset Loader.
    Args:
        trainset: training set
        valset: validation set
        batch size
        shuffle ratio
        num_workers: number of workers in DataLoader

    Returns:
        trainloader (torch.utils.data.DataLoader): DataLoader for training set
        testloader (torch.utils.data.DataLoader): DataLoader for validation set
    """
    transformations = transforms.Compose(
        [transforms.Lambda(lambda x: (x / 127.5) - 1.0)])

    # Load training data and validation data
    training_set = TripletDataset(dataroot,trainset, transformations)
    trainloader = DataLoader(training_set,
                             batch_size=batch_size,
                             shuffle=shuffle,
                             num_workers=num_workers)

    validation_set = TripletDataset(dataroot, valset, transformations)
    valloader = DataLoader(validation_set,
                           batch_size=batch_size,
                           shuffle=shuffle,
                           num_workers=num_workers)

    return trainloader, valloader

# Define model
class NetworkNvidia(nn.Module):
    """NVIDIA model used in the paper."""

    def __init__(self):
        """The NVIDIA architecture.
            Data preprocessing and image normalisation
            Convolution: 5x5, filter: 24, strides: 2x2, activation: ELU
            The 5 convolution layers are for feature extraction.
            The fully connected layers are predict the steering angless

        """
        super(NetworkNvidia, self).__init__()
        self.conv_layers = nn.Sequential(
            # convolution layer 1
            nn.Conv2d(3, 24, 5, stride=2),
            nn.ELU(),
            # convolution layer 2
            nn.Conv2d(24, 36, 5, stride=2),
            nn.ELU(),
            # convolution layer 3
            nn.Conv2d(36, 48, 5, stride=2),
            nn.ELU(),
            # convolution layer 4
            nn.Conv2d(48, 64, 3),
            nn.ELU(),
            # convolution layer 5
            nn.Conv2d(64, 64, 3),
            nn.Dropout(0.5)
        )
        self.linear_layers = nn.Sequential(
            # fully connected layer 1
            nn.Linear(in_features=64 * 2 * 425, out_features=150),
            nn.ELU(),
            # fully connected layer 2
            nn.Linear(in_features=150, out_features=80),
            nn.ELU(),
            # fully connected layer 3
            nn.Linear(in_features=80, out_features=10),
            # fully connected layer 4
            nn.Linear(in_features=10, out_features=1)
        )

    def forward(self, input):
        """Forward propogation"""
        # change the tensor shape
        input = input.view(input.size(0), 3, 196, 455)
        # pass the input to the 5 convolution layers
        output = self.conv_layers(input)
        # reshape the features to pass it into activation layers
        output = output.view(output.size(0), -1)
        # pass the feature vectors to the fully connected layers
        output = self.linear_layers(output)
        return output

class Trainer(object):
    """Trainer."""

    def __init__(self,ckptroot,model,device,epochs,criterion,optimizer,scheduler,start_epoch,trainloader,validationloader):
        """Self-Driving car Trainer.
        Args:
            model:
            device:
            epochs:
            criterion:
            optimizer:
            start_epoch:
            trainloader:
            validationloader:

        """
        super(Trainer, self).__init__()

        self.model = model
        self.device = device
        self.epochs = epochs
        self.ckptroot = ckptroot
        self.criterion = criterion
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.start_epoch = start_epoch
        self.trainloader = trainloader
        self.validationloader = validationloader

    def train(self):
        """Training process."""
        self.model.to(self.device)
        for epoch in range(self.start_epoch, self.epochs + self.start_epoch):
            self.scheduler.step()

            # Training
            train_loss = 0.0
            train_acc = 0.0
            self.model.train()

            for local_batch, (centers) in enumerate(self.trainloader):
                # Transfer to GPU
                centers= toDevice(centers, self.device)
                # Model computations
                self.optimizer.zero_grad()
                datas = [centers]

                for data in datas:
                    imgs, angles = data
                    outputs = self.model(imgs)
                    #print (outputs)
                    loss = self.criterion(outputs, angles.unsqueeze(1))
                    correct = (angles.eq(outputs.long())).sum()
                    # back propagation
                    loss.backward()
                    self.optimizer.step()
                    #calculate the traimning loss
                    train_loss += loss.data.item()
                    train_acc += correct.data.item()

                if local_batch % 500 == 0:

                    print("Training Epoch: {} | Loss: {}".format(epoch, train_loss / (local_batch + 1)))
                    print("Training Accuracy: {} | Acc: {}".format(epoch, (train_acc / (local_batch + 1))))
            # Validation
            self.model.eval()
            valid_loss = 0
            valid_acc = 0
            with torch.set_grad_enabled(False):
                for local_batch, (centers) in enumerate(self.validationloader):
                    # Transfer to GPU
                    centers = toDevice(centers, self.device)

                    # Model computations
                    self.optimizer.zero_grad()
                    datas = [centers]
                    for data in datas:
                        # extract the dataloader
                        imgs, angles = data
                        # model prediction
                        outputs = self.model(imgs)
                        # calculating the model loss
                        loss = self.criterion(outputs, angles.unsqueeze(1))
                        correct = (angles.eq(outputs.long())).sum()

                        valid_loss += loss.data.item()
                        valid_acc += correct.data.item()
                    if local_batch % 500 == 0:
                        print("Validation Loss: {}".format(valid_loss / (local_batch + 1)))
                        print("Accuracy: {}".format((valid_acc / (local_batch + 1))))

            print()
            # Save model
            if epoch % 5 == 0 or epoch == self.epochs + self.start_epoch - 1:
                state = {
                    'epoch': epoch + 1,
                    'state_dict': self.model.state_dict(),
                    'optimizer': self.optimizer.state_dict(),
                    'scheduler': self.scheduler.state_dict(),
                }

                self.save_checkpoint(state)

    def save_checkpoint(self, state):
        """Save checkpoint."""
        print("==> Save checkpoint ...")
        if not os.path.exists(self.ckptroot):
            os.makedirs(self.ckptroot)
        # save the trained model
        torch.save(state, self.ckptroot + 'real.ckpt')

def main():
    print(torch.cuda.device_count())
    print(torch.cuda.is_available())

    ## Load the data you want ot test
    data_path = "training_set.txt"
    dataroot = []
    steers = []
    ## data loading
    with open(data_path) as file:
        for line in file:
            if line.split(',')[0]=="center":continue
            dataroot.append('driving_dataset/' + line.split(' ')[0])
            steers.append(line.split(' ')[1].strip())
    # model hyper parameters
    lr = 1e-5
    weight_decay = 1e-5
    batch_size = 5 ##32
    num_workers = 8
    test_size = 0.8
    shuffle = True
    epochs = 1 #80
    start_epoch = 0
    save = 'new_training_model/'
    trainset, valset = load_data(dataroot,steers, test_size)
    trainloader, validationloader = data_loader(dataroot,
                                            trainset, valset,
                                            batch_size,
                                            shuffle,
                                            num_workers)
    # Define model
    print("==> Initialize model ...")
    model = NetworkNvidia()
    print("==> Initialize model done ...")
    # Define optimizer and criterion
    optimizer = optim.Adam(model.parameters(),
                        lr=lr,
                        weight_decay=weight_decay)
    # Define the loss function
    criterion = nn.MSELoss()
    # learning rate scheduler
    scheduler = MultiStepLR(optimizer, milestones=[30, 50], gamma=0.1)
    # Use Gpu if available..else run CPU
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Accesing you device CPU/GPU?(cuda)",device)
    print("==> Start training ...")
    # starting the imitation learning training
    trainer = Trainer(save,
                    model,
                    device,
                    epochs,
                    criterion,
                    optimizer,
                    scheduler,
                    start_epoch,
                    trainloader,
                    validationloader)
    trainer.train()

if __name__ == "__main__":
    main()


import unittest
import random
import numpy
import copy
import time

from aupyom import Sound, Sampler
from aupyom.util import example_audio_file

class TestTimeStretch(unittest.TestCase):
    def setUp(self):
        self.s1 = Sound.from_file(example_audio_file())

        freq, dt, sr = 440.0, 10.0, 22050
        t = numpy.linspace(0, dt, sr * dt)
        self.s2 = Sound(numpy.sin(2 * numpy.pi * freq * t), sr)

        self.sounds = [self.s1, self.s2, copy.deepcopy(self.s1)]

        self.sampler = Sampler(backend='dummy')

    def test_sampling(self):
        for s in self.sounds:
            self.sampler.play(s)
            time.sleep(random.random())
            s.pitch_shift = random.random() * 24 - 12
            time.sleep(random.random())
            s.stretch_factor = random.random() * 5

        while self.sampler.sounds:
            time.sleep(.1)

if __name__ == '__main__':
    unittest.main()