from railrl.data_management.env_replay_buffer import EnvReplayBuffer
import railrl.torch.pytorch_util as ptu
import railrl.misc.hyperparameter as hyp

from railrl.samplers.data_collector import MdpPathCollector
from railrl.torch.sac.policies import GaussianPolicy, GaussianMixturePolicy
from railrl.launchers.launcher_util import run_experiment
from railrl.torch.sac.policies import (
    MakeDeterministic, TanhGaussianPolicy,
)
from railrl.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.awr_sac import AWRSACTrainer
from railrl.demos.source.mdp_path_loader import MDPPathLoader
from railrl.demos.source.dict_to_mdp_path_loader import DictToMDPPathLoader
from railrl.launchers.arglauncher import run_variants

import roboverse

import numpy as np
import os.path as osp
from railrl.core import logger

import torch

DEFAULT_BUFFER = ('/media/avi/data/Work/github/jannerm/bullet-manipulation/data'
                  '/feb25_SawyerReach-v0_state_2K_dense_reward_randomize_noise'
                  '_std_0.2/reach_demos_state_noise02_vector.npy')

def experiment(variant):
    expl_env = roboverse.make(variant['env'], gui=False, randomize=True,
                              observation_mode='state', reward_type='shaped',
                              transpose_image=True)
    eval_env = expl_env
    action_dim = int(np.prod(eval_env.action_space.shape))
    obs_dim = 11

    M = variant['layer_size']
    qf1 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    qf2 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_qf1 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    target_qf2 = FlattenMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_sizes=[M, M],
    )
    policy_class = variant.get("policy_class", TanhGaussianPolicy)
    policy = policy_class(
        obs_dim=obs_dim,
        action_dim=action_dim,
        **variant['policy_kwargs'],
    )

    buffer_policy = policy_class(
        obs_dim=obs_dim,
        action_dim=action_dim,
        **variant['policy_kwargs'],
    )

    trainer = AWRSACTrainer(
        env=eval_env,
        policy=policy,
        qf1=qf1,
        qf2=qf2,
        target_qf1=target_qf1,
        target_qf2=target_qf2,
        buffer_policy=buffer_policy,
        **variant['trainer_kwargs']
    )

    expl_policy = policy
    expl_path_collector = MdpPathCollector(
        expl_env,
        expl_policy,
    )
    eval_policy = MakeDeterministic(policy)
    eval_path_collector = MdpPathCollector(
        eval_env,
        eval_policy,
    )

    replay_buffer_kwargs = dict(
        max_replay_buffer_size=variant['replay_buffer_size'],
        env=expl_env,
    )

    replay_buffer = EnvReplayBuffer(**replay_buffer_kwargs)

    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=expl_env,
        evaluation_env=eval_env,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        max_path_length=variant['max_path_length'],
        batch_size=variant['batch_size'],
        num_epochs=variant['num_epochs'],
        num_eval_steps_per_epoch=variant['num_eval_steps_per_epoch'],
        num_expl_steps_per_train_loop=variant['num_expl_steps_per_train_loop'],
        num_trains_per_train_loop=variant['num_trains_per_train_loop'],
        min_num_steps_before_training=variant['min_num_steps_before_training'],
    )

    algorithm.to(ptu.device)

    demo_train_buffer = EnvReplayBuffer(
        **replay_buffer_kwargs,
    )
    demo_test_buffer = EnvReplayBuffer(
        **replay_buffer_kwargs,
    )

    path_loader_kwargs = variant.get("path_loader_kwargs", {})

    save_paths = None  # FIXME(avi)
    if variant.get('save_paths', False):
        algorithm.post_train_funcs.append(save_paths)
    if variant.get('load_demos', False):
        path_loader_class = variant.get('path_loader_class', MDPPathLoader)
        path_loader = path_loader_class(trainer,
            replay_buffer=replay_buffer,
            demo_train_buffer=demo_train_buffer,
            demo_test_buffer=demo_test_buffer,
            **path_loader_kwargs
        )
        path_loader.load_demos()
    if variant.get('pretrain_policy', False):
        trainer.pretrain_policy_with_bc()
    if variant.get('pretrain_rl', False):
        trainer.pretrain_q_with_bc_data()
    if variant.get('save_pretrained_algorithm', False):
        p_path = osp.join(logger.get_snapshot_dir(), 'pretrain_algorithm.p')
        pt_path = osp.join(logger.get_snapshot_dir(), 'pretrain_algorithm.pt')
        data = algorithm._get_snapshot()
        data['algorithm'] = algorithm
        torch.save(data, open(pt_path, "wb"))
        torch.save(data, open(p_path, "wb"))
    if variant.get('train_rl', True):
        algorithm.train()

if __name__ == "__main__":

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--env", type=str, default='SawyerReach-v0',
                        choices=('SawyerReach-v0', 'SawyerGraspOne-v0'))
    parser.add_argument("--buffer", type=str, default=DEFAULT_BUFFER)
    parser.add_argument("--gpu", type=int, default=1)

    args = parser.parse_args()

    variant = dict(
        num_epochs=1001,
        num_eval_steps_per_epoch=250,
        num_trains_per_train_loop=1000,
        num_expl_steps_per_train_loop=1000,
        min_num_steps_before_training=1000,
        max_path_length=50,
        batch_size=1024,
        replay_buffer_size=int(1E6),

        layer_size=256,
        policy_class=GaussianMixturePolicy,
        policy_kwargs=dict(
            hidden_sizes=[256, 256, 256, 256],
            max_log_std=0,
            min_log_std=-6,
            std_architecture="values",
            num_gaussians=1,
        ),

        algorithm="SAC",
        version="normal",
        collection_mode='batch',
        trainer_kwargs=dict(
            discount=0.99,
            soft_target_tau=5e-3,
            target_update_period=1,
            policy_lr=3E-4,
            qf_lr=3E-4,
            reward_scale=1,
            beta=0.1,
            use_automatic_entropy_tuning=False,
            alpha=0,
            compute_bc=False,

            bc_num_pretrain_steps=0,
            q_num_pretrain1_steps=0,
            q_num_pretrain2_steps=25000,
            policy_weight_decay=1e-4,
            q_weight_decay=0,
            bc_loss_type="mle",
            awr_loss_type="mle",

            rl_weight=1.0,
            use_awr_update=True,
            use_reparam_update=False,
            reparam_weight=0.0,
            awr_weight=1.0,
            bc_weight=1.0,

            post_bc_pretrain_hyperparams=dict(
                bc_weight=0.0,
                compute_bc=False,
            ),

            reward_transform_kwargs=None, # r' = r + 1
            terminal_transform_kwargs=None, # t = 0
        ),
        num_exps_per_instance=1,
        region='us-west-2',

        path_loader_class=DictToMDPPathLoader,
        path_loader_kwargs=dict(
            obs_key="state_observation",
            demo_paths=[
                dict(
                    path=args.buffer,
                    obs_dict=False,
                    is_demo=True,
                ),
                # dict(
                #     path="demos/icml2020/hand/pen_bc5.npy",
                #     obs_dict=False,
                #     is_demo=False,
                #     train_split=0.9,
                # ),
            ],
        ),
        add_env_demos=True,
        add_env_offpolicy_data=True,

        # logger_variant=dict(
        #     tensorboard=True,
        # ),
        load_demos=True,
        pretrain_policy=True,
        pretrain_rl=True,
        # save_pretrained_algorithm=True,
        # snapshot_mode="all",
    )

    variant['env'] = args.env

    search_space = {
        # 'env': ["pen-sparse-v0", "door-sparse-v0", "relocate-sparse-v0", ],
        'trainer_kwargs.bc_loss_type': ["mle"],
        'trainer_kwargs.awr_loss_type': ["mle"],
        'seedid': [1],
        'trainer_kwargs.beta': [0.1, 0.3, 0.5],
        'trainer_kwargs.reparam_weight': [0.0, ],
        'trainer_kwargs.awr_weight': [1.0],
        'trainer_kwargs.bc_weight': [1.0, ],
        'policy_kwargs.std_architecture': ["values", ],

        # 'trainer_kwargs.compute_bc': [True, ],
        'trainer_kwargs.awr_use_mle_for_vf': [False, ],
        'trainer_kwargs.awr_sample_actions': [False, ],
        'trainer_kwargs.awr_min_q': [True, ],

        # 'trainer_kwargs.q_weight_decay': [1e-5, 1e-4, ],

        'trainer_kwargs.reward_transform_kwargs': [None, ],
        'trainer_kwargs.terminal_transform_kwargs': [dict(m=0, b=0), ],
        # 'policy_kwargs.num_gaussians': [1, 2, ],
    }

    sweeper = hyp.DeterministicHyperparameterSweeper(
        search_space, default_parameters=variant,
    )

    mode = 'local'
    exp_prefix = 'railrl-AWRSAC-{}-state'.format(args.env)
    n_seeds = 1

    from joblib import Parallel, delayed
    variants = []
    for variant in sweeper.iterate_hyperparameters():
        variants.append(variant)
    num_threads = len(variants)
    Parallel(n_jobs=num_threads)(delayed(run_experiment)(
        experiment,
        exp_name=exp_prefix,
        mode=mode,
        variant=variant,
        use_gpu=True,
        gpu_id=args.gpu,
        unpack_variant=False,
    ) for variant in variants)


import numpy as np
from typing import Callable, Type, NamedTuple, Optional
import copy
import logging
import ConfigSpace as CS

from .debug_log import DebugLogPrinter
from .gp_profiling import GPMXNetSimpleProfiler
from .hp_ranges import HyperparameterRanges_CS
from ..datatypes.common import CandidateEvaluation, Candidate, candidate_for_print, PendingEvaluation
from ..datatypes.hp_ranges import HyperparameterRanges
from ..datatypes.tuning_job_state import TuningJobState
from ..models.gp_model import GaussProcSurrogateModel, GPModel
from ..models.gpmodel_skipopt import SkipOptimizationPredicate
from ..models.gpmodel_transformers import GPModelPendingCandidateStateTransformer, GPModelArgs
from ..tuning_algorithms.base_classes import LocalOptimizer, AcquisitionFunction, ScoringFunction
from ..tuning_algorithms.bo_algorithm import BayesianOptimizationAlgorithm
from ..tuning_algorithms.bo_algorithm_components import IndependentThompsonSampling
from ..tuning_algorithms.common import RandomStatefulCandidateGenerator, compute_blacklisted_candidates
from ..tuning_algorithms.default_algorithm import dictionarize_objective, \
    DEFAULT_METRIC, DEFAULT_LOCAL_OPTIMIZER_CLASS, \
    DEFAULT_NUM_INITIAL_CANDIDATES, DEFAULT_NUM_INITIAL_RANDOM_EVALUATIONS
from ..utils.duplicate_detector import DuplicateDetectorIdentical

logger = logging.getLogger(__name__)

GET_CONFIG_RANDOM_RETRIES = 50

class MapReward(NamedTuple):
    forward: Callable[[float], float]
    reverse: Callable[[float], float]

    def __call__(self, x: float) -> float:
        return self.forward(x)

def accumulate_profiling_record(
        cum_records: dict, profiler: GPMXNetSimpleProfiler, pick_random: bool):
    # Pull out profiling data from this block
    block_id = profiler.id_counter - 1
    curr_record = {
        r.tag: r.duration for r in profiler.records
        if r.id == block_id}
    if pick_random:
        curr_record['num_random'] = 1
        curr_record['total_all'] = curr_record['random']
    else:
        curr_record['num_model'] = 1
    # Sum up entries
    for k, v in curr_record.items():
        if k in cum_records:
            cum_records[k] += v
        else:
            cum_records[k] = v

SUPPORTED_INITIAL_SCORING = {
    'thompson_indep',
    'acq_func'}

DEFAULT_INITIAL_SCORING = 'thompson_indep'

def create_initial_candidates_scorer(
        initial_scoring: str, model: GaussProcSurrogateModel,
        acquisition_class: Type[AcquisitionFunction],
        random_state: np.random.RandomState) -> ScoringFunction:
    if initial_scoring == 'thompson_indep':
        return IndependentThompsonSampling(model, random_state=random_state)
    else:
        return acquisition_class(model)

def check_initial_candidates_scorer(initial_scoring: str) -> str:
    if initial_scoring is None:
        return DEFAULT_INITIAL_SCORING
    else:
        assert initial_scoring in SUPPORTED_INITIAL_SCORING, \
            "initial_scoring = '{}' is not supported".format(
                initial_scoring)
        return initial_scoring

class GPFIFOSearcher(object):
    """
    Supports standard GP-based hyperparameter optimization, when used with a
    FIFO scheduler.

    """
    def __init__(
            self, hp_ranges: HyperparameterRanges, random_seed: int,
            gpmodel: GPModel, model_args: GPModelArgs,
            map_reward: MapReward,
            acquisition_class: Type[AcquisitionFunction],
            init_state: TuningJobState = None,
            local_minimizer_class: Type[LocalOptimizer] = DEFAULT_LOCAL_OPTIMIZER_CLASS,
            skip_optimization: SkipOptimizationPredicate = None,
            num_initial_candidates: int = DEFAULT_NUM_INITIAL_CANDIDATES,
            num_initial_random_choices: int = DEFAULT_NUM_INITIAL_RANDOM_EVALUATIONS,
            initial_scoring: Optional[str] = None,
            profiler: Optional[GPMXNetSimpleProfiler] = None,
            first_is_default: bool = True,
            debug_log: Optional[DebugLogPrinter] = None,
            cost_metric_name: Optional[str] = None):
        """
        Note that the SurrogateModel is created on demand (by the state
        transformer) in get_config, along with components needed for the BO
        algorithm.

        The searcher is supposed to maximize reward, while internally, the
        criterion is minimized. map_reward maps reward to internal criterion, it
        must be strictly decreasing.

        :param hp_ranges: Configuration space without resource attribute
        :param random_seed:
        :param gpmodel: GP regression model
        :param model_args: Arguments for GPMXNet model creation
        :param map_reward: Function mapping reward to criterion to be minimized
        :param acquisition_class: Type for acquisition function
        :param init_state: TuningJobState to start from (default is empty)
        :param local_minimizer_class: Type for local minimizer
        :param skip_optimization: Predicate, see
            GPMXNetPendingCandidateStateTransformer
        :param num_initial_candidates: See BayesianOptimizationAlgorithm
        :param num_initial_random_choices: Configs are sampled at random until
            this many candidates received label feedback
        :param initial_scoring: Scoring function to rank initial candidates.
            Default: thompson_indep (independent Thompson sampling)
        :param profiler: If given, HPO computations are profiled
        :param first_is_default: If true, the first config to be evaluated is
            the default of the search space. Otherwise, the first is sampled
            at random
        :param debug_log: DebugLogPrinter for debug logging (optional)

        """
        self.hp_ranges = hp_ranges
        self.random_seed = random_seed
        self.num_initial_candidates = num_initial_candidates
        self.num_initial_random_choices = num_initial_random_choices
        self.map_reward = map_reward
        self.local_minimizer_class = local_minimizer_class
        self.acquisition_class = acquisition_class
        self.debug_log = debug_log
        self.initial_scoring = check_initial_candidates_scorer(initial_scoring)
        # Create state transformer
        # Initial state is empty (note that the state is mutable)
        if init_state is None:
            init_state = TuningJobState(
                hp_ranges=hp_ranges,
                candidate_evaluations=[],
                failed_candidates=[],
                pending_evaluations=[])
        else:
            assert hp_ranges is init_state.hp_ranges, \
                "hp_ranges and init_state.hp_ranges must be same object"
        self.state_transformer = GPModelPendingCandidateStateTransformer(
            gpmodel=gpmodel,
            init_state=init_state,
            model_args=model_args,
            skip_optimization=skip_optimization,
            profiler=profiler,
            debug_log=debug_log)
        self.random_state = np.random.RandomState(random_seed)
        self.random_generator = RandomStatefulCandidateGenerator(
            hp_ranges, random_state=self.random_state)
        self.profiler = profiler
        self.do_profile = (profiler is not None)
        self.first_is_default = first_is_default
        if first_is_default:
            assert isinstance(hp_ranges, HyperparameterRanges_CS), \
                "If first_is_default, must have hp_ranges of HyperparameterRanges_CS type"
        if debug_log is not None:
            assert isinstance(hp_ranges, HyperparameterRanges_CS), \
                "If debug_log is given, must have hp_ranges of HyperparameterRanges_CS type"
        if cost_metric_name is not None:
            self.cost_metric_name = cost_metric_name
        else:
            self.cost_metric_name = 'elapsed_time'
        # Sums up profiling records across all get_config calls
        self._profile_record = dict()
        if debug_log is not None:
            deb_msg = "[GPFIFOSearcher.__init__]\n"
            deb_msg += ("- acquisition_class = {}\n".format(acquisition_class))
            deb_msg += ("- local_minimizer_class = {}\n".format(local_minimizer_class))
            deb_msg += ("- num_initial_candidates = {}\n".format(num_initial_candidates))
            deb_msg += ("- num_initial_random_choices = {}\n".format(num_initial_random_choices))
            deb_msg += ("- initial_scoring = {}\n".format(self.initial_scoring))
            deb_msg += ("- first_is_default = {}".format(first_is_default))
            logger.info(deb_msg)

    def update(self, config: Candidate, reward: float, **kwargs):
        """
        Registers new datapoint at config, with reward reward.
        Note that in general, config should previously have been registered as
        pending (register_pending). If so, it is switched from pending
        to labeled. If not, it is considered directly labeled.

        :param config:
        :param reward:
        """
        crit_val = self.map_reward(reward)
        metrics = dictionarize_objective(crit_val)
        if 'elapsed_time' in kwargs:
            metrics[self.cost_metric_name] = kwargs['elapsed_time']
        self.state_transformer.label_candidate(CandidateEvaluation(
            candidate=copy.deepcopy(config), metrics=metrics))
        if self.debug_log is not None:
            config_id = self.debug_log.config_id(config)
            msg = "Update for config_id {}: reward = {}, crit_val = {}".format(
                config_id, reward, crit_val)
            logger.info(msg)

    def register_pending(self, config: Candidate):
        """
        Registers config as pending. This means the corresponding evaluation
        task is running. Once it finishes, update is called for config.

        """
        # It is OK for the candidate already to be registered as pending, in
        # which case we do nothing
        state = self.state_transformer.state
        if config not in state.pending_candidates:
            if config in (x.candidate for x in state.candidate_evaluations):
                evals = state.candidate_evaluations
                num_labeled = len(evals)
                pos_cand = next(
                    i for i, x in enumerate(evals) if x.candidate == config)
                error_msg = """
                This configuration is already registered as labeled:
                   Position of labeled candidate: {} of {}
                   Label value: {}
                """.format(
                    pos_cand, num_labeled,
                    evals[pos_cand].metrics[DEFAULT_METRIC])
                assert False, error_msg
            self.state_transformer.append_candidate(config)

    def get_config(self) -> Candidate:
        """
        Runs Bayesian optimization in order to suggest the next config to evaluate.

        :return: Next config to evaluate at
        """
        state = self.state_transformer.state
        if self.do_profile:
            fit_hyperparams = not self.state_transformer.skip_optimization(
                state)
            self.profiler.set_state(state, fit_hyperparams)
        blacklisted_candidates = compute_blacklisted_candidates(state)
        pick_random = (len(blacklisted_candidates) < self.num_initial_random_choices) or \
            (not state.candidate_evaluations)
        if self.debug_log is not None:
            self.debug_log.start_get_config('random' if pick_random else 'BO')
        if pick_random:
            config = None
            if self.first_is_default and (not blacklisted_candidates):
                # Use default configuration if there is one specified
                default_config = self.hp_ranges.config_space.get_default_configuration()
                if default_config and len(default_config.get_dictionary()) > 0:
                    config = default_config
                    if self.debug_log is not None:
                        logger.info("Start with default config:\n{}".format(
                            candidate_for_print(config)))
            if config is None:
                if self.do_profile:
                    self.profiler.start('random')
                for _ in range(GET_CONFIG_RANDOM_RETRIES):
                    _config = self.hp_ranges.random_candidate(self.random_state)
                    if _config not in blacklisted_candidates:
                        config = _config
                        break
                if config is None:
                    raise AssertionError(
                        "Failed to sample a configuration not already chosen "
                        "before. Maybe there are no free configurations left? "
                        "The blacklist size is {}".format(len(blacklisted_candidates)))
                if self.do_profile:
                    self.profiler.stop('random')
        else:
            # Obtain current SurrogateModel from state transformer. Based on
            # this, the BO algorithm components can be constructed
            state = self.state_transformer.state
            if self.do_profile:
                self.profiler.start('total_all')
                self.profiler.start('total_update')
            # Note: Asking for the model triggers the posterior computation
            model = self.state_transformer.model()
            if self.do_profile:
                self.profiler.stop('total_update')
            # Create BO algorithm
            initial_candidates_scorer = create_initial_candidates_scorer(
                self.initial_scoring, model, self.acquisition_class,
                self.random_state)
            local_optimizer = self.local_minimizer_class(
                state, model, self.acquisition_class)
            # Make sure not to use the same random seed for each call:
            #random_seed = compute_random_seed({'0': state}, self.random_seed)
            bo_algorithm = BayesianOptimizationAlgorithm(
                initial_candidates_generator=self.random_generator,
                initial_candidates_scorer=initial_candidates_scorer,
                num_initial_candidates=self.num_initial_candidates,
                local_optimizer=local_optimizer,
                pending_candidate_state_transformer=None,
                blacklisted_candidates=blacklisted_candidates,
                num_requested_candidates=1,
                greedy_batch_selection=False,
                duplicate_detector=DuplicateDetectorIdentical(),
                profiler=self.profiler,
                sample_unique_candidates=False,
                debug_log=self.debug_log)
            # Next candidate decision
            if self.do_profile:
                self.profiler.start('total_nextcand')
            _config = bo_algorithm.next_candidates()
            if len(_config) == 0:
                raise AssertionError(
                    "Failed to find a configuration not already chosen "
                    "before. Maybe there are no free configurations left? "
                    "The blacklist size is {}".format(len(blacklisted_candidates)))
            config = _config[0]
            if self.do_profile:
                self.profiler.stop('total_nextcand')
            if self.do_profile:
                self.profiler.stop('total_all')

        if self.debug_log is not None:
            self.debug_log.set_final_config(config)
            # All get_config debug log info is only written here
            self.debug_log.write_block()
        if self.do_profile:
            self.profiler.clear()
            # Pull out profiling data from this block, add to _profile_record
            accumulate_profiling_record(
                self._profile_record, self.profiler, pick_random)

        return config

    def evaluation_failed(self, config: Candidate):
        # Remove pending candidate
        self.state_transformer.drop_candidate(config)
        # Mark config as failed (which means it will be blacklisted in
        # future get_config calls)
        self.state_transformer.mark_candidate_failed(config)

    def dataset_size(self):
        return len(self.state_transformer.state.candidate_evaluations)

    def cumulative_profile_record(self):
        """
        If profiling is activated, we sum up the profiling blocks for each
        call of get_config and return it as dict. See get_config for what
        is recorded:
        - num_random: Number of get_config calls with random selection
        - num_model: Number of get_config calls with model-based selection
        - total_all: Sum of total times for all get_config calls
        """
        return self._profile_record

    def get_params(self):
        """
        Note: Once MCMC is supported, this method will have to be refactored.

        :return: Dictionary with current hyperparameter values
        """
        return self.state_transformer.get_params()

    def set_params(self, param_dict):
        self.state_transformer.set_params(param_dict)

    def get_state(self):
        """
        The mutable state consists of the GP model parameters, the
        TuningJobState, and the skip_optimization predicate (which can have a
        mutable state).
        We assume that skip_optimization can be pickled.

        """
        state = {
            'model_params': self.get_params(),
            'state': encode_state(self.state_transformer.state),
            'skip_optimization': self.state_transformer.skip_optimization,
            'random_state': self.random_state}
        if self.debug_log is not None:
            state['debug_log'] = self.debug_log.get_mutable_state()
        return state

    def clone_from_state(self, state):
        # Create clone with mutable state taken from 'state'
        model_args = self.state_transformer._model_args
        init_state = decode_state(state['state'], self.hp_ranges)
        skip_optimization = state['skip_optimization']
        new_searcher = GPFIFOSearcher(
            hp_ranges=self.hp_ranges,
            random_seed=self.random_seed,
            gpmodel=self.state_transformer._gpmodel,
            model_args=model_args,
            map_reward=self.map_reward,
            acquisition_class=self.acquisition_class,
            init_state=init_state,
            local_minimizer_class=self.local_minimizer_class,
            skip_optimization=skip_optimization,
            num_initial_candidates=self.num_initial_candidates,
            num_initial_random_choices=self.num_initial_random_choices,
            initial_scoring=self.initial_scoring,
            profiler=self.profiler,
            first_is_default=self.first_is_default,
            debug_log=self.debug_log)
        new_searcher.state_transformer.set_params(state['model_params'])
        new_searcher.random_state = state['random_state']
        new_searcher.random_generator.random_state = \
            state['random_state']
        if self.debug_log and 'debug_log' in state:
            new_searcher.debug_log.set_mutable_state(state['debug_log'])
        # Invalidate self (must not be used afterwards)
        self.state_transformer = None
        return new_searcher

def encode_state(state: TuningJobState) -> dict:
    assert isinstance(state.hp_ranges, HyperparameterRanges_CS), \
        "Must have hp_ranges of HyperparameterRanges_CS type"
    candidate_evaluations = [
        {'candidate': eval.candidate.get_dictionary(),
         'metrics': eval.metrics}
        for eval in state.candidate_evaluations]
    failed_candidates = [x.get_dictionary() for x in state.failed_candidates]
    pending_evaluations = [
        eval.candidate.get_dictionary() for eval in state.pending_evaluations]
    return {
        'candidate_evaluations': candidate_evaluations,
        'failed_candidates': failed_candidates,
        'pending_evaluations': pending_evaluations}

def decode_state(enc_state: dict, hp_ranges: HyperparameterRanges_CS) \
        -> TuningJobState:
    assert isinstance(hp_ranges, HyperparameterRanges_CS), \
        "Must have hp_ranges of HyperparameterRanges_CS type"
    config_space = hp_ranges.config_space

    def to_cs(x):
        return CS.Configuration(config_space, values=x)

    candidate_evaluations = [
        CandidateEvaluation(to_cs(x['candidate']), x['metrics'])
        for x in enc_state['candidate_evaluations']]
    failed_candidates = [to_cs(x) for x in enc_state['failed_candidates']]
    pending_evaluations = [
        PendingEvaluation(to_cs(x)) for x in enc_state['pending_evaluations']]
    return TuningJobState(
        hp_ranges=hp_ranges,
        candidate_evaluations=candidate_evaluations,
        failed_candidates=failed_candidates,
        pending_evaluations=pending_evaluations)

def map_reward(const=1.0) -> MapReward:
    """
    Factory for map_reward argument in GPMultiFidelitySearcher.
    """
    def const_minus_x(x):
        return const - x

    return MapReward(forward=const_minus_x, reverse=const_minus_x)


import os
import math
import random
import numpy as np
from tqdm import tqdm
from mmsdk import mmdatasdk
from sklearn.metrics import f1_score, accuracy_score
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.tensorboard import SummaryWriter

torch.backends.cudnn.deterministic = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.backends.cudnn.benchmark = True

# parameters
user = 'dango/multimodal/CMU-MOSEI'
data_dir = '/home/'+user+'/align/'
log_dir = '/home/'+user+'/par_log/log_fix_out2/'
L_DIM = 300
V_DIM = 35
A_DIM = 74
L_LEN = 50
V_LEN = 50
A_LEN = 50
CLIP = 1.0
EPOCHS = 99
BATCH = 64
DIM = 96
N_HEADS = 6
FFN = 2
N_LAYERS = 2
LR = 0.001
DROP = 0.0
P_LEN = 6

# data
data_dict = {'linguistic':data_dir+'glove_vectors.csd', 'acoustic':data_dir+'COAVAREP.csd', 'visual':data_dir+'FACET 4.2.csd', 'label':data_dir+'All Labels.csd'}
data_set = mmdatasdk.mmdataset(data_dict)

train_name, test_name = [], []
for name in data_set.computational_sequences['visual'].data.keys():
    if name.split('[')[0] in mmdatasdk.cmu_mosei.standard_folds.standard_test_fold:
        test_name.append(name.split('[')[0])
    else:
        train_name.append(name.split('[')[0])
train_name, test_name = set(train_name), set(test_name)

def name_list(name):
    name_list = []
    for n in name:
        temp_list = []
        count = 0
        for i in range((98//P_LEN+1)*P_LEN):
            if n+'['+str(i)+']' in data_set.computational_sequences['visual'].data.keys():
                temp_list.append(n+'['+str(i)+']')
            else:
                temp_list.append('no_name')
            count += 1
            if count == P_LEN:
                count = 0
                if temp_list[0] != 'no_name':
                    name_list.append(temp_list)
                temp_list = []
    return name_list
train_name_list = name_list(train_name)
test_name_list = name_list(test_name)

def masking(m, m_len):
    if len(m) >= m_len:
        m_mask = np.ones(m_len)
    else:
        m_mask = np.concatenate((np.ones(len(m)), np.zeros(m_len - len(m))))
    m = np.concatenate([m, np.zeros([m_len]+list(m.shape[1:]))],axis=0)[:m_len,...]
    for i in range(len(m)):
        for j in range(len(m[i])):
            if math.isinf(m[i][j]) or math.isnan(m[i][j]):
                m[i][j] = -71.
    return m, m_mask

def label_processing(l):
    label = l[1:]
    label[0] = 1 if label[0] > 0 else 0
    label[1] = 1 if label[1] > 0 else 0
    label[2] = 1 if label[2] > 0 else 0
    label[3] = 1 if label[3] > 0 else 0
    label[4] = 1 if label[4] > 0 else 0
    label[5] = 1 if label[5] > 0 else 0
    return label

def data_loader(data_set, name_list, batch_size):
    random.shuffle(name_list)
    count = 0
    while count < len(name_list):
        batch = []
        size = min(batch_size, len(name_list) - count)
        for _ in range(size):
            l_temp, v_temp, a_temp, lmask_temp, vmask_temp, amask_temp, label_temp, mask_temp = [], [], [], [], [], [], [], []
            for i in range(P_LEN):
                if name_list[count][i] != 'no_name':
                    l, l_mask = masking(data_set.computational_sequences['linguistic'].data[name_list[count][i]]["features"][-L_LEN:], L_LEN)
                    v, v_mask = masking(data_set.computational_sequences['visual'].data[name_list[count][i]]["features"][-V_LEN:], V_LEN)
                    a, a_mask = masking(data_set.computational_sequences['acoustic'].data[name_list[count][i]]["features"][-A_LEN:], A_LEN)
                    label = label_processing(data_set.computational_sequences['label'].data[name_list[count][i]]["features"][0])
                    mask = 1
                else:
                    l, v, a = np.zeros((L_LEN, L_DIM)), np.zeros((V_LEN, V_DIM)), np.zeros((A_LEN, A_DIM))
                    l_mask, v_mask, a_mask = np.zeros(L_LEN), np.zeros(V_LEN), np.zeros(A_LEN)
                    label = np.array([0, 0, 0, 0, 0, 0])
                    mask = 0
                l_temp.append(np.expand_dims(l, axis=0))
                v_temp.append(np.expand_dims(v, axis=0))
                a_temp.append(np.expand_dims(a, axis=0))
                lmask_temp.append(np.expand_dims(l_mask, axis=0))
                vmask_temp.append(np.expand_dims(v_mask, axis=0))
                amask_temp.append(np.expand_dims(a_mask, axis=0))
                label_temp.append(np.expand_dims(label, axis=0))
                mask_temp.append(mask)
            batch.append((np.concatenate(l_temp, axis=0), np.concatenate(v_temp, axis=0), np.concatenate(a_temp, axis=0), np.concatenate(label_temp, axis=0),\
                          np.concatenate(lmask_temp, axis=0), np.concatenate(vmask_temp, axis=0), np.concatenate(amask_temp, axis=0), np.asarray(mask_temp)))
            count += 1
        yield batch

# model
def get_parameter_number(net):
    total_num = sum(p.numel() for p in net.parameters())
    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}

class Unify_Dimension_Conv1d(nn.Module):
    def __init__(self, l_dim, v_dim, a_dim, dim):
        super().__init__()
        self.linguistic = nn.Conv1d(l_dim, dim, kernel_size=1, bias=False)
        self.visual = nn.Conv1d(v_dim, dim, kernel_size=1, bias=False)
        self.acoustic = nn.Conv1d(a_dim, dim, kernel_size=1, bias=False)
        self.drop = nn.Dropout(DROP)
    def forward(self, l, v, a):
        l, v, a = l.transpose(1, 2), v.transpose(1, 2), a.transpose(1, 2)
        l, v, a = self.drop(self.linguistic(l)), self.drop(self.visual(v)), self.drop(self.acoustic(a))
        return l.transpose(1, 2), v.transpose(1, 2), a.transpose(1, 2)

class Position_Embedding(nn.Module):
    def __init__(self, max_len, dim):
        super().__init__()
        self.position_embeddings = nn.Embedding(max_len, dim)
        self.len = max_len
    def forward(self, x):
        position_ids = torch.arange(self.len, device=device).unsqueeze(0).repeat(x.size()[0],1)
        return self.position_embeddings(position_ids.to(device))

class Attention_Block(nn.Module):
    def __init__(self, dim, n_heads):
        super().__init__()
        self.w_qkv = nn.ModuleList([nn.Linear(dim, dim, bias = False) for _ in range(3)])
        self.n_heads = n_heads
        self.drop = nn.Dropout(DROP)
        self.proj = nn.Linear(dim, dim, bias = False)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, FFN * dim),
            nn.ReLU(),
            nn.Linear(FFN * dim, dim),
            nn.Dropout(DROP),
        )
        self.a = nn.Parameter(torch.FloatTensor([0]), requires_grad=True)
        self.b = nn.Parameter(torch.FloatTensor([0]), requires_grad=True)
        self.c = nn.Parameter(torch.FloatTensor([0]), requires_grad=True)
    def split_last(self, x, shape):
        shape = list(shape)
        assert shape.count(-1) <= 1
        if -1 in shape:
            shape[shape.index(-1)] = int(x.size(-1) / -np.prod(shape))
        return x.view(*x.size()[:-1], *shape)
    def merge_last(self, x, n_dims):
        s = x.size()
        assert n_dims > 1 and n_dims < len(s)
        return x.view(*s[:-n_dims], -1)
    def multi_head_attention(self, q, k, v, mask, scores = None):
        '''
        q    → (batch_size, q_len, dim);
        k, v → (batch_size, kv_len, dim);
        mask → (batch_size, kv_len) / (batch_size, q_len, kv_len)
        '''
        q, k, v = self.w_qkv[0](q), self.w_qkv[1](k), self.w_qkv[2](v)
        q, k, v = (self.split_last(x, (self.n_heads, -1)).transpose(1, 2) for x in [q, k, v])
        if scores is not None:
            scores = q @ k.transpose(-2, -1) / np.sqrt(k.size(-1)) + self.c * scores
        else:
            scores = q @ k.transpose(-2, -1) / np.sqrt(k.size(-1))
        if mask is not None:
            if len(mask.shape) == 2:
                mask = mask[:, None, None, :]
            elif len(mask.shape) == 3:
                mask = torch.unsqueeze(mask, 1)
                mask = mask.repeat(1,self.n_heads,1,1)
            scores -= 1.0e8 * (1.0 - mask)
        att = F.softmax(scores, dim=-1)
        q = (att @ v).transpose(1, 2).contiguous()
        q = self.merge_last(q, 2)
        return self.drop(self.proj(q)), scores
    def forward(self, q, k, v, mask, scores = None):
        x, scores =  self.multi_head_attention(q, k, v, mask, scores)
        q = self.norm1(q + self.a * x)
        q = self.norm2(q + self.b * self.ffn(q))
        return q, scores

class Multi_class(nn.Module):
    def __init__(self, l_dim, v_dim, a_dim, dim, l_len, v_len, a_len, n_heads, n_layers, ffn):
        super().__init__()
        self.unify_dimension = Unify_Dimension_Conv1d(l_dim, v_dim, a_dim, dim)
        self.linguistic_position = Position_Embedding(l_len, dim)
        self.visual_position = Position_Embedding(v_len, dim)
        self.acoustic_position = Position_Embedding(a_len, dim)
        self.n_layers = n_layers
        self.multimodal_blocks = nn.ModuleList([Attention_Block(dim, n_heads) for _ in range(9*n_layers)])
        self.fully_connected = nn.Linear(dim*6, dim)
        self.normalization = nn.LayerNorm(dim)
        self.drop = nn.Dropout(DROP)
    def forward(self, l, v, a, l_mask, v_mask, a_mask):
        l, v, a = self.unify_dimension(l, v, a)
        l = l + self.linguistic_position(l)
        v = v + self.visual_position(v)
        a = a + self.acoustic_position(a)
        ll, lv, la = l, l, l
        vv, vl, va = v, v, v
        aa, al, av = a, a, a
        scores = None
        for i in range(self.n_layers):
            ll, scores = self.multimodal_blocks[self.n_layers*0+i](ll, l, l, l_mask, scores)
        scores = None
        for i in range(self.n_layers):
            lv, scores = self.multimodal_blocks[self.n_layers*1+i](lv, v, v, v_mask, scores)
        scores = None
        for i in range(self.n_layers):
            la, scores = self.multimodal_blocks[self.n_layers*2+i](la, a, a, a_mask, scores)
        scores = None
        for i in range(self.n_layers):
            vv, scores = self.multimodal_blocks[self.n_layers*3+i](vv, v, v, v_mask, scores)
        scores = None
        for i in range(self.n_layers):
            vl, scores = self.multimodal_blocks[self.n_layers*4+i](vl, l, l, l_mask, scores)
        scores = None
        for i in range(self.n_layers):
            va, scores = self.multimodal_blocks[self.n_layers*5+i](va, a, a, a_mask, scores)
        scores = None
        for i in range(self.n_layers):
            aa, scores = self.multimodal_blocks[self.n_layers*6+i](aa, a, a, a_mask, scores)
        scores = None
        for i in range(self.n_layers):
            al, scores = self.multimodal_blocks[self.n_layers*7+i](al, l, l, l_mask, scores)
        scores = None
        for i in range(self.n_layers):
            av, scores = self.multimodal_blocks[self.n_layers*8+i](av, v, v, v_mask, scores)
        l = torch.cat([ll, lv, la], dim=2)
        v = torch.cat([vv, vl, va], dim=2)
        a = torch.cat([aa, al, av], dim=2)
        x = torch.cat([l, a, v], dim=1)  # (batch, l_len+v_len+a_len, dim*3)
        x = torch.cat([torch.mean(x, 1), torch.max(x, 1)[0]], dim=1)  # (batch, dim*6)
        x = self.drop(nn.ReLU()(self.normalization(self.fully_connected(x))))  # (batch, dim)
        return x

class State_Transfer(nn.Module):
    def __init__(self, l_dim, v_dim, a_dim, dim, l_len, v_len, a_len, n_heads, n_layers, ffn):
        super().__init__()
        self.feature = Multi_class(l_dim=l_dim, v_dim=v_dim, a_dim=a_dim, dim=dim, l_len=l_len, v_len=v_len, a_len=a_len, n_heads=n_heads, n_layers=n_layers, ffn=ffn)
        self.classifier = nn.Linear(dim, 6*2)
        self.trans = nn.Parameter(torch.rand(6,6), requires_grad=True)
    def forward(self, l, v, a, l_mask, v_mask, a_mask):  # (batch, y_len, x_len, x_dim)
        out_list, feats_list = [], []
        for i in range(l.shape[1]):
            temp_l, temp_v, temp_a, temp_l_mask, temp_v_mask, temp_a_mask = l[:,i,:,:], v[:,i,:,:], a[:,i,:,:], l_mask[:,i,:], v_mask[:,i,:], a_mask[:,i,:]
            temp_feats = self.feature(temp_l, temp_v, temp_a, temp_l_mask, temp_v_mask, temp_a_mask)  # (batch, dim)
            temp_feats = self.classifier(temp_feats)  # (batch, 6*2)
            temp_out_t1, temp_feats = temp_feats.chunk(2, 1)  # (batch, 6)
            if i != 0:
                alpha = torch.sigmoid(temp_feats + feats_list[-1].squeeze())
                temp_out_t0 = torch.tanh(torch.matmul(out_list[-1].squeeze(), self.trans))
                temp_out_t1 = (1-alpha) * temp_out_t1 + alpha * temp_out_t0
            out_list.append(temp_out_t1.unsqueeze(1))
            feats_list.append(temp_feats.unsqueeze(1))
        out = torch.cat(out_list, dim=1)
        return out  # (batch, y_len, 6)

# run
def multi_circle_loss(y_pred, y_true):
    y_pred = (1 - 2 * y_true) * y_pred
    y_pred_neg = y_pred - y_true * 1e12
    y_pred_pos = y_pred - (1 - y_true) * 1e12
    zeros = torch.zeros_like(y_pred[..., :1], dtype = torch.float)
    y_pred_neg = torch.cat([y_pred_neg, zeros], dim=-1)
    y_pred_pos = torch.cat([y_pred_pos, zeros], dim=-1)
    neg_loss = torch.logsumexp(y_pred_neg, dim=-1)
    pos_loss = torch.logsumexp(y_pred_pos, dim=-1)
    return neg_loss + pos_loss

def train(model, iterator, optimizer):
    model.train()
    epoch_loss, count = 0, 0
    iter_bar = tqdm(iterator, desc='Training')
    for _, batch in enumerate(iter_bar):
        count += 1
        optimizer.zero_grad()
        linguistic, visual, acoustic, label, l_mask, v_mask, a_mask, mask = zip(*batch)
        linguistic, visual, acoustic, label, l_mask, v_mask, a_mask, mask = torch.cuda.FloatTensor(linguistic), torch.cuda.FloatTensor(visual), torch.cuda.FloatTensor(acoustic),\
        torch.cuda.LongTensor(label), torch.cuda.FloatTensor(l_mask), torch.cuda.FloatTensor(v_mask), torch.cuda.FloatTensor(a_mask), torch.cuda.LongTensor(mask)
        logits_clsf = model(linguistic, visual, acoustic, l_mask, v_mask, a_mask)
        loss = multi_circle_loss(logits_clsf, label)
        loss = (loss*mask).mean()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), CLIP)  #梯度裁剪
        optimizer.step()
        iter_bar.set_description('Iter (loss=%3.3f)'%loss.item())
        epoch_loss += loss.item()
    return epoch_loss / count

def valid(model, iterator):
    model.eval()
    epoch_loss, count = 0, 0
    with torch.no_grad():
        iter_bar = tqdm(iterator, desc='Validation')
        for _, batch in enumerate(iter_bar):
            count += 1
            linguistic, visual, acoustic, label, l_mask, v_mask, a_mask, mask = zip(*batch)
            linguistic, visual, acoustic, label, l_mask, v_mask, a_mask, mask = torch.cuda.FloatTensor(linguistic), torch.cuda.FloatTensor(visual), torch.cuda.FloatTensor(acoustic),\
            torch.cuda.LongTensor(label), torch.cuda.FloatTensor(l_mask), torch.cuda.FloatTensor(v_mask), torch.cuda.FloatTensor(a_mask), torch.cuda.LongTensor(mask)
            logits_clsf = model(linguistic, visual, acoustic, l_mask, v_mask, a_mask)
            loss = multi_circle_loss(logits_clsf, label)
            loss = (loss*mask).mean()
            iter_bar.set_description('Iter (loss=%3.3f)'%loss.item())
            epoch_loss += loss.item()
    return epoch_loss, count, epoch_loss / count

def run(model, data_set, train_list, valid_list, batch_size, learning_rate, epochs, name):
    log_file = log_dir+name+'.txt'
    with open(log_file, 'w') as log_f:
        log_f.write('epoch, train_loss, valid_loss\n')
    writer = SummaryWriter(log_dir)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=2, verbose=True)
    stop = 0
    loss_list = []
    for epoch in range(epochs):
        print('Epoch: ' + str(epoch+1))
        train_iterator = data_loader(data_set, train_list, batch_size)
        valid_iterator = data_loader(data_set, valid_list, batch_size)
        train_loss = train(model, train_iterator, optimizer)
        _, _, valid_loss = valid(model, valid_iterator)
        writer.add_scalars(name, {'train_loss':train_loss, 'valid_loss':valid_loss}, epoch)
        scheduler.step(valid_loss)
        loss_list.append(valid_loss)
        with open(log_file, 'a') as log_f:
            log_f.write('\n{epoch},{train_loss: 2.2f},{valid_loss: 2.2f}\n'.format(epoch=epoch+1, train_loss=train_loss, valid_loss=valid_loss))
        if valid_loss == min(loss_list):
            stop = 0
            torch.save(model.state_dict(), os.path.join(log_dir, name+'_'+str(valid_loss)[:4]+'.pt'))
        else:
            stop += 1
            if stop >= 4:
                break
    writer.close()

random.shuffle(train_name_list)
model_1 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
valid_list = train_name_list[:int(len(train_name_list)*0.2)]
train_list = train_name_list[int(len(train_name_list)*0.2):]
run(model_1, data_set, train_list, valid_list, batch_size=BATCH, learning_rate=LR, epochs=EPOCHS, name='model_1')
model_2 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
valid_list = train_name_list[int(len(train_name_list)*0.2):int(len(train_name_list)*0.4)]
train_list = train_name_list[:int(len(train_name_list)*0.2)] + train_name_list[int(len(train_name_list)*0.4):]
run(model_2, data_set, train_list, valid_list, batch_size=BATCH, learning_rate=LR, epochs=EPOCHS, name='model_2')
model_3 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
valid_list = train_name_list[int(len(train_name_list)*0.4):int(len(train_name_list)*0.6)]
train_list = train_name_list[:int(len(train_name_list)*0.4)] + train_name_list[int(len(train_name_list)*0.6):]
run(model_3, data_set, train_list, valid_list, batch_size=BATCH, learning_rate=LR, epochs=EPOCHS, name='model_3')
model_4 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
valid_list = train_name_list[int(len(train_name_list)*0.6):int(len(train_name_list)*0.8)]
train_list = train_name_list[:int(len(train_name_list)*0.6)] + train_name_list[int(len(train_name_list)*0.8):]
run(model_4, data_set, train_list, valid_list, batch_size=BATCH, learning_rate=LR, epochs=EPOCHS, name='model_4')
model_5 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
valid_list = train_name_list[int(len(train_name_list)*0.8):]
train_list = train_name_list[:int(len(train_name_list)*0.8)]
run(model_5, data_set, train_list, valid_list, batch_size=BATCH, learning_rate=LR, epochs=EPOCHS, name='model_5')

# for name, p in model_1.named_parameters():
#     if name == 'trans.weight':
#         print(p)
# for name, p in model_2.named_parameters():
#     if name == 'trans.weight':
#         print(p)

def test(model_1, model_2):
    model_1.eval()
    model_2.eval()
    label_happ, soft_happ = [], []
    label_sadn, soft_sadn = [], []
    label_ange, soft_ange = [], []
    label_surp, soft_surp = [], []
    label_disg, soft_disg = [], []
    label_fear, soft_fear = [], []
    best_happ_f1, best_happ_acc, best_happ_t = 0, 0, 0
    best_sadn_f1, best_sadn_acc, best_sadn_t = 0, 0, 0
    best_ange_f1, best_ange_acc, best_ange_t = 0, 0, 0
    best_surp_f1, best_surp_acc, best_surp_t = 0, 0, 0
    best_disg_f1, best_disg_acc, best_disg_t = 0, 0, 0
    best_fear_f1, best_fear_acc, best_fear_t = 0, 0, 0
    with torch.no_grad():
        for t in tqdm(range(400)):
            threshold = t/200-1.0
            test_iterator = data_loader(data_set, test_name_list, batch_size=BATCH)
            for _, batch in enumerate(test_iterator):
                linguistic, visual, acoustic, label, l_mask, v_mask, a_mask, mask = zip(*batch)
                linguistic, visual, acoustic, label, l_mask, v_mask, a_mask, mask = torch.cuda.FloatTensor(linguistic), torch.cuda.FloatTensor(visual), torch.cuda.FloatTensor(acoustic),\
                torch.cuda.LongTensor(label), torch.cuda.FloatTensor(l_mask), torch.cuda.FloatTensor(v_mask), torch.cuda.FloatTensor(a_mask), torch.cuda.LongTensor(mask)
                pred_1 = (model_1(linguistic, visual, acoustic, l_mask, v_mask, a_mask)).cpu().detach()
                pred_2 = (model_2(linguistic, visual, acoustic, l_mask, v_mask, a_mask)).cpu().detach()
                pred = pred_1 * 0.6 + pred_2 * 0.4
                zero = torch.zeros_like(pred)
                one = torch.ones_like(pred)
                pred = torch.where(pred > threshold, one, zero)
                label = label.cpu().detach()
                for i in range(len(mask)):
                    for j in range(P_LEN):
                        if int(mask[i][j]) == 1:
                            label_happ.append(int(label[i][j][0]))
                            label_sadn.append(int(label[i][j][1]))
                            label_ange.append(int(label[i][j][2]))
                            label_surp.append(int(label[i][j][3]))
                            label_disg.append(int(label[i][j][4]))
                            label_fear.append(int(label[i][j][5]))
                            soft_happ.append(int(pred[i][j][0]))
                            soft_sadn.append(int(pred[i][j][1]))
                            soft_ange.append(int(pred[i][j][2]))
                            soft_surp.append(int(pred[i][j][3]))
                            soft_disg.append(int(pred[i][j][4]))
                            soft_fear.append(int(pred[i][j][5]))
                        else:
                            break
            happ_acc = accuracy_score(label_happ, soft_happ)
            happ_f1 = f1_score(label_happ, soft_happ, average='weighted')
            sadn_acc = accuracy_score(label_sadn, soft_sadn)
            sadn_f1 = f1_score(label_sadn, soft_sadn, average='weighted')
            ange_acc = accuracy_score(label_ange, soft_ange)
            ange_f1 = f1_score(label_ange, soft_ange, average='weighted')
            surp_acc = accuracy_score(label_surp, soft_surp)
            surp_f1 = f1_score(label_surp, soft_surp, average='weighted')
            disg_acc = accuracy_score(label_disg, soft_disg)
            disg_f1 = f1_score(label_disg, soft_disg, average='weighted')
            fear_acc = accuracy_score(label_fear, soft_fear)
            fear_f1 = f1_score(label_fear, soft_fear, average='weighted')
            if happ_f1 > best_happ_f1:
                best_happ_f1 = happ_f1
                best_happ_acc = happ_acc
                best_happ_t = threshold
            if sadn_f1 > best_sadn_f1:
                best_sadn_f1 = sadn_f1
                best_sadn_acc = sadn_acc
                best_sadn_t = threshold
            if ange_f1 > best_ange_f1:
                best_ange_f1 = ange_f1
                best_ange_acc = ange_acc
                best_ange_t = threshold
            if surp_f1 > best_surp_f1:
                best_surp_f1 = surp_f1
                best_surp_acc = surp_acc
                best_surp_t = threshold
            if disg_f1 > best_disg_f1:
                best_disg_f1 = disg_f1
                best_disg_acc = disg_acc
                best_disg_t = threshold
            if fear_f1 > best_fear_f1:
                best_fear_f1 = fear_f1
                best_fear_acc = fear_acc
                best_fear_t = threshold
    return best_happ_f1, best_happ_acc, best_happ_t, best_sadn_f1, best_sadn_acc, best_sadn_t, best_ange_f1, best_ange_acc, best_ange_t, best_surp_f1, best_surp_acc, best_surp_t,\
    best_disg_f1, best_disg_acc, best_disg_t, best_fear_f1, best_fear_acc, best_fear_t

# log_fix_out1
model_1 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
model_1.load_state_dict(torch.load(log_dir + 'model_2_1.33.pt'))
model_2 = State_Transfer(l_dim=L_DIM, v_dim=V_DIM, a_dim=A_DIM, dim=DIM, l_len=L_LEN, v_len=V_LEN, a_len=A_LEN, n_heads=N_HEADS, n_layers=N_LAYERS, ffn=FFN).to(device)
model_2.load_state_dict(torch.load(log_dir + 'model_1_1.37.pt'))

best_happ_f1, best_happ_acc, best_happ_t, best_sadn_f1, best_sadn_acc, best_sadn_t, best_ange_f1, best_ange_acc, best_ange_t, best_surp_f1, best_surp_acc, best_surp_t,\
best_disg_f1, best_disg_acc, best_disg_t, best_fear_f1, best_fear_acc, best_fear_t = test(model_1, model_2)
print('best_happ_acc: ', best_happ_acc)
print('best_happ_f1: ', best_happ_f1)
print('best_happ_t: ', best_happ_t)
print('best_sadn_acc: ', best_sadn_acc)
print('best_sadn_f1: ', best_sadn_f1)
print('best_sadn_t: ', best_sadn_t)
print('best_ange_acc: ', best_ange_acc)
print('best_ange_f1: ', best_ange_f1)
print('best_ange_t: ', best_ange_t)
print('best_fear_acc: ', best_fear_acc)
print('best_fear_f1: ', best_fear_f1)
print('best_fear_t: ', best_fear_t)
print('best_disg_acc: ', best_disg_acc)
print('best_disg_f1: ', best_disg_f1)
print('best_disg_t: ', best_disg_t)
print('best_surp_acc: ', best_surp_acc)
print('best_surp_f1: ', best_surp_f1)
print('best_surp_t: ', best_surp_t)


import numpy as np

class Agent(object):
    def __init__(self, dim_action):
        self.dim_action = dim_action

    def act(self, ob, reward, done):
        return np.tanh(np.random.randn(self.dim_action)) # random action


import numpy
import cv2
import argparse
import librosa
import gc
import os
import multiprocessing
from moviepy.editor import CompositeAudioClip, AudioFileClip, VideoFileClip, VideoClip
from tempfile import gettempdir
from headbang.motion import OpenposeDetector, bpm_from_beats, align_beats_motion
from headbang.params import DEFAULTS
from headbang import HeadbangBeatTracker
from headbang.util import load_wav, overlay_clicks
from madmom.io.audio import write_wave_file

def main():
    parser = argparse.ArgumentParser(
        description="Track human pose in videos with music alongside groove metrics and beat tracking"
    )
    parser.add_argument(
        "--keypoints",
        type=str,
        default=DEFAULTS["pose_keypoints"],
        help="Override the default face keypoints (default=%(default)s)",
    )
    parser.add_argument(
        "--bpm-history",
        type=float,
        default=DEFAULTS["bpm_history"],
        help="History of video (in seconds) to be included in the window of current bpm computation (default=%(default)s)",
    )
    parser.add_argument(
        "--event-threshold-frames",
        type=int,
        default=DEFAULTS["event_thresh_frames"],
        help="Threshold in number of frames by which an event is considered to be the same (default=%(default)s)",
    )
    parser.add_argument(
        "--debug-motion",
        action="store_true",
        help="Only perform motion detection with matplotlib with debug plot",
    )
    parser.add_argument(
        "--debug-bpm",
        action="store_true",
        help="Use with --debug-motion, also display bpm estimation",
    )
    parser.add_argument(
        "--experimental-wav-out",
        type=str,
        default="",
        help="wav output path for bop clicks",
    )
    parser.add_argument(
        "--experimental-bop-align",
        type=float,
        default=DEFAULTS["bop_align"],
        help="align bops and beats within this duration window (s) (default=%(default)s)",
    )
    parser.add_argument(
        "--experimental-sick-chain-boundary",
        type=float,
        default=DEFAULTS["sick_chain_boundary"],
        help="time boundary to separate sick portions of a song (default=%(default)s)",
    )
    parser.add_argument(
        "--experimental-sick-chain",
        action="store_true",
        help="identify and display when a sick chain is occuring (according to the boundary argument)",
    )

    parser.add_argument("mp4_in", type=str, help="mp4 file to process")
    parser.add_argument("mp4_out", type=str, help="mp4 output path")

    args = parser.parse_args()

    video_path = args.mp4_in
    cap = cv2.VideoCapture(video_path)

    if args.debug_motion:
        strong_beat_locations = numpy.empty(shape=(1,), dtype=numpy.float32)
        all_beat_locations = numpy.empty(shape=(1,), dtype=numpy.float32)
    else:
        pool = multiprocessing.Pool(multiprocessing.cpu_count() - 1)
        hbt = HeadbangBeatTracker(pool)

        audio, _ = librosa.load(video_path, sr=44100, dtype=numpy.float32, mono=True)

        # get beat locations
        print("Getting beat locations using consensus beat tracking")
        strong_beat_locations = hbt.beats(audio)

        # pre onset alignment
        all_beat_locations = hbt.beat_consensus

    total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)

    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_duration = 1 / fps

    pose_tracker = OpenposeDetector(
        total_frames,
        frame_duration,
        keypoints=args.keypoints,
    )

    total_duration = frame_duration * total_frames

    video_width = 1920
    video_height = 1080

    bop_pos = (100, int(video_height - 50))
    bop_bpm_pos = (280, int(video_height - 50))
    bop_color = (255, 255, 0)

    all_beat_pos = (635, int(video_height - 50))
    all_beat_bpm_pos = (860, int(video_height - 50))
    all_beat_color = (255, 0, 0)

    groove_pos = (1200, int(video_height - 50))
    groove_color = (0, 255, 0)

    strong_beat_pos = (1600, int(video_height - 50))
    strong_beat_color = (0, 165, 255)

    sick_color = (255, 0, 255)

    alpha = 0.90

    blank_frame = numpy.zeros((video_height, video_width, 3), numpy.uint8)

    def process_first_pass(*args, **kwargs):
        grabbed, frame = cap.read()
        if not grabbed:
            return blank_frame

        frame = cv2.resize(frame, (video_width, video_height))

        # update latest pose frame
        out_frame = pose_tracker.process_frame(frame)

        # draw semi-transparent rectangle
        hud_overlay = out_frame.copy()

        cv2.rectangle(
            hud_overlay,
            (0, video_height - 150),
            (video_width, video_height),
            (0, 0, 0),
            -1,
        )
        video_frame = cv2.addWeighted(hud_overlay, alpha, out_frame, 1 - alpha, 0)

        # adjust color on frames
        video_frame = cv2.cvtColor(video_frame, cv2.COLOR_BGR2RGB)
        return video_frame

    tmp_mp4 = os.path.join(gettempdir(), "headbang-hud-tmp.mp4")

    print("Processing video - first pass with pose detection")
    out_clip = VideoClip(make_frame=process_first_pass, duration=total_duration)
    out_clip.write_videofile(tmp_mp4, fps=fps)

    # get head bop locations by indexing into time array
    all_time = numpy.linspace(0, frame_duration * total_frames, int(total_frames))

    # take top peaks only
    print("Getting peaks of y motion")
    peaks = pose_tracker.find_peaks()
    bop_locations = all_time[peaks]

    aligned_bop_locations = numpy.asarray(
        align_beats_motion(
            all_beat_locations, bop_locations, args.experimental_bop_align
        )
    )

    event_thresh = args.event_threshold_frames * frame_duration

    print("Marking beat and head bop positions on output frames")

    print("run a gc, just in case...")
    gc.collect()

    all_beats_bpm = 0
    bop_bpm = 0
    time_since_last_groove = None

    bop_bpm_plot_history = []

    # define a function to filter the first video to add more stuff
    def process_second_pass(get_frame_fn, frame_time):
        nonlocal all_beats_bpm, bop_bpm, time_since_last_groove, bop_bpm_plot_history
        frame = get_frame_fn(frame_time)

        frame_max = frame_time
        frame_min = max(0, frame_time - args.bpm_history)

        all_beat_history = all_beat_locations[
            numpy.where(
                (all_beat_locations >= frame_min) & (all_beat_locations <= frame_max)
            )
        ]

        # keep a running history of bops
        bop_history = bop_locations[
            numpy.where((bop_locations >= frame_min) & (bop_locations <= frame_max))
        ]

        all_beats_bpm_tmp = bpm_from_beats(all_beat_history)
        bop_bpm_tmp = bpm_from_beats(bop_history)

        if not numpy.isnan(all_beats_bpm_tmp):
            all_beats_bpm = all_beats_bpm_tmp

        if not numpy.isnan(bop_bpm_tmp):
            bop_bpm = bop_bpm_tmp

        if args.debug_bpm:
            bop_bpm_plot_history.append((bop_history, bop_bpm))

        is_strong_beat = False
        is_beat = False
        is_bop = False
        is_sick = False
        if any(
            [b for b in all_beat_locations if numpy.abs(b - frame_time) <= event_thresh]
        ):
            is_beat = True
        if any(
            [
                b
                for b in strong_beat_locations
                if numpy.abs(b - frame_time) <= event_thresh
            ]
        ):
            is_strong_beat = True
        if any([b for b in bop_locations if numpy.abs(b - frame_time) <= event_thresh]):
            is_bop = True

        is_groove = is_bop and is_beat

        if (
            time_since_last_groove is not None
            and frame_time - time_since_last_groove
            <= args.experimental_sick_chain_boundary
        ):
            is_sick = True
        else:
            # sick chain is broken
            time_since_last_groove = None

        if is_groove and time_since_last_groove is None:
            time_since_last_groove = frame_time

        if not args.debug_motion:
            if is_beat:
                cv2.putText(
                    frame,
                    "BEAT",
                    all_beat_pos,
                    cv2.FONT_HERSHEY_SIMPLEX,
                    2.0,
                    all_beat_color,
                    3,
                    cv2.LINE_AA,
                )
            cv2.putText(
                frame,
                "{0:.2f} bpm".format(all_beats_bpm),
                all_beat_bpm_pos,
                cv2.FONT_HERSHEY_SIMPLEX,
                1.0,
                all_beat_color,
                2,
                cv2.LINE_AA,
            )

            if is_strong_beat:
                cv2.putText(
                    frame,
                    "BEAT+",
                    strong_beat_pos,
                    cv2.FONT_HERSHEY_SIMPLEX,
                    2.0,
                    strong_beat_color,
                    3,
                    cv2.LINE_AA,
                )
            if is_groove:
                cv2.putText(
                    frame,
                    "GROOVE",
                    groove_pos,
                    cv2.FONT_HERSHEY_SIMPLEX,
                    2.0,
                    groove_color,
                    3,
                    cv2.LINE_AA,
                )

        if is_bop:
            cv2.putText(
                frame,
                "BOP",
                bop_pos,
                cv2.FONT_HERSHEY_SIMPLEX,
                2.0,
                bop_color,
                3,
                cv2.LINE_AA,
            )
        cv2.putText(
            frame,
            "{0:.2f} bpm".format(bop_bpm),
            bop_bpm_pos,
            cv2.FONT_HERSHEY_SIMPLEX,
            1.0,
            bop_color,
            2,
            cv2.LINE_AA,
        )

        if args.experimental_sick_chain and is_sick:
            font = cv2.FONT_HERSHEY_SIMPLEX
            text = "SICK"

            textsize = cv2.getTextSize(text, font, 26, 12)[0]

            x_text = int((frame.shape[1] - textsize[0]) / 2)
            y_text = int((frame.shape[0] + textsize[1]) / 2) - 75

            cv2.putText(
                frame, text, (x_text, y_text), font, 26, sick_color, 12, cv2.LINE_AA
            )

        return frame

    print(
        "Processing video - second pass with marked event locations and estimated tempos"
    )

    print("run a gc, just in case...")
    gc.collect()

    out_clip_tmp = VideoFileClip(tmp_mp4, fps_source="fps")
    out_clip2 = out_clip_tmp.fl(process_second_pass)

    audio_clip = AudioFileClip(video_path)
    new_audioclip = CompositeAudioClip([audio_clip])

    out_clip2.audio = new_audioclip
    out_clip2.write_videofile(args.mp4_out, fps=fps)

    print("cleaning up tmp mp4")
    os.remove(tmp_mp4)

    if args.experimental_wav_out:
        print("Overlaying clicks at bop locations")
        x_stereo = load_wav(video_path, stereo=True)
        x_with_clicks = overlay_clicks(x_stereo, aligned_bop_locations)

        print("Writing output with clicks to {0}".format(args.experimental_wav_out))
        write_wave_file(x_with_clicks, args.experimental_wav_out, sample_rate=44100)

    if args.debug_motion:
        print("Displaying debug y coordinate and bpm plots")
        pose_tracker.plot_ycoords(bop_bpm_plot_history, args.debug_bpm)


from flask import Flask, send_from_directory
from flask_cors import CORS
import threading
import logging
from sys import argv
import sys
import utils
import constants
from collections import deque
from io import BytesIO
from wave import open as wave_open
from pyaudio import PyAudio
from math import ceil, floor
from time import sleep, perf_counter
from tensorflow import lite
from os import path as os_path
from numpy import argmax, int16, frombuffer, ndarray
from numpy import array as np_array
import webbrowser
from itertools import islice

direction = "n/a"

'''
------------------------------------------------------------------------------------------------------------------------------

HTTP SERVER CODE

------------------------------------------------------------------------------------------------------------------------------
'''

def resource_path(relative_path):
    try:
        base_path = sys._MEIPASS
    except Exception:
        base_path = os_path.abspath(".")

    return os_path.join(base_path, relative_path)

log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

def run_http_server():
    global direction
    app = Flask(__name__, static_folder=resource_path('./dist'))
    CORS(app)

    # return the last inferred direction
    @app.route("/")
    def main():
        global direction
        return direction

    # Serve React App
    @app.route('/app', defaults={'path': ''})
    @app.route('/app/<path:path>')
    def serve(path):
        print(path)
        if path != "" and os_path.exists(app.static_folder + '/' + path):
            return send_from_directory(app.static_folder, path)
        else:
            return send_from_directory(app.static_folder, 'index.html')

    # serve the flask application
    app.run(
        host=constants.webserver_host_name,
        port=constants.webserver_port,
        debug=True,
        use_reloader=False
    )

'''
------------------------------------------------------------------------------------------------------------------------------

MACHINE LEARNING CODE

------------------------------------------------------------------------------------------------------------------------------
'''

model_a = model_b = input_details = output_details = None

commands_a = ['eight', 'zero', 'right', 'down', 'left', 'two', '_background_noise_', 'stop', 'go', 'up']
commands_b = ['bed', 'right', 'down', 'left', '_background_noise_', 'no', 'wow', 'up', 'yes', 'five']

# Initialize a Tensorflow Lite model from a provided path
def init_tflite_model(path):
    tflite_model = lite.Interpreter(model_path=path)
    tflite_model.allocate_tensors()
    input_details = tflite_model.get_input_details()
    output_details = tflite_model.get_output_details()

    return tflite_model, input_details, output_details

# Classify the provided audio data using the provided model and return
# the result if the constraints are met
def run_model(inference_array, model, commands, results, index):
    model.set_tensor(input_details[0]['index'], inference_array)
    model.invoke()

    result = model.get_tensor(output_details[0]['index'])
    prediction = argmax(result, axis=1)

    thresh_adjust = 0.0
    if commands[int(prediction[0])] == "up":
        thresh_adjust = 0.45

    if result[0][prediction] > (0.45 + thresh_adjust):
        direction = commands[int(prediction[0])]
        results[index] = direction
    else:
        results[index] = "n/a"

# Infer the intended keyword from a provided audio binary
def execute_models(audio_binary):
    # print('Audio binary received. Starting inference..')
    spectrogram = utils.binary_to_spec(audio_binary)

    inference_array = []
    inference_array.append(spectrogram.numpy())
    inference_array = np_array(inference_array)

    results = {
        'a': None,
        'b': None
    }

    thread_a = threading.Thread(
        target=run_model,
        args=(inference_array, model_a, commands_a, results, 'a'),
        name="inference_a"
    )

    thread_b = threading.Thread(
        target=run_model,
        args=(inference_array, model_b, commands_b, results, 'b'),
        name="inference_b"
    )

    # We chose a threaded model for the inference step because the
    # total classification time decreased from ~4ms to ~2ms
    thread_a.start()
    thread_b.start()
    thread_a.join()
    thread_b.join()

    result_a = results['a']
    result_b = results['b']

    prediction = result_a if result_a == result_b else 'n/a'
    return prediction

# Stream audio data from the selected microphone into a 1 second buffer
def listen(device_index, buffer, shared):
    p = PyAudio()

    stream = p.open(
        format = constants.FORMAT,
        channels = constants.CHANNELS,
        rate = constants.RATE,
        input = True,
        output = False,
        frames_per_buffer = constants.CHUNK,
        input_device_index=device_index
    )

    predict_start = perf_counter()

    while True:
        current_time = perf_counter()
        data = stream.read(constants.CHUNK, exception_on_overflow = False)
        buffer.append(data)

        if (len(buffer) < ceil(constants.RATE / constants.CHUNK)):
            continue

        data = frombuffer(data, dtype=int16)
        max = ndarray.max(data)
        if (max > 600): # check if current chunk surpasses energy threshold
            predict_start = perf_counter()
            shared['predict'] = True
        if (current_time - predict_start > 1.5): # classify for 1.5 seconds
            shared['predict'] = False

        buffer.popleft()

def predict(buffer, shared):
    global direction
    predictions = deque()

    while True:
        sleep(0.003)
        if (shared['predict'] == False):
            continue
        sleep(0.003)
        wav_data = None
        with BytesIO() as wav_file:
            wav_writer = wave_open(wav_file, "wb")
            try:
                wav_writer.setframerate(constants.RATE)
                wav_writer.setsampwidth(2)
                wav_writer.setnchannels(1)
                wav_writer.writeframes(b"".join(islice(buffer, 0, floor((constants.RATE / constants.CHUNK) / 0.6))))
                wav_data = wav_file.getvalue()
            finally:  # make sure resources are cleaned up
                wav_writer.close()

        predition = execute_models(wav_data)
        predictions.append(predition)
        if (len(predictions) > 2):
            predictions.popleft()

        # Only update direction if the previous two preditions are the same
        if len(predictions) == 2 and predictions[0] == predictions[1] and direction != predictions[0]:
            direction = predictions[0]
            if predictions[0] != 'n/a': print(f'predition: {direction}')

# Listen and infer keywords from the selected microphone
def run_inference():
    p = PyAudio()

    info = p.get_host_api_info_by_index(0)
    numdevices = info.get('deviceCount')
    for i in range(0, numdevices):
            if (p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
                print("Input Device id ", i, " - ", p.get_device_info_by_host_api_device_index(0, i).get('name'))

    device_index = 0

    if (len(argv) > 1):
      device_index = int(argv[1])
    else:
      print('\nEnter the device id you would like to use:')
      device_index = int(input())

    buffer = deque()
    shared = {
        'predict': False
    }

    print('Starting listening thread...')
    threading.Thread(
        target=lambda: listen(device_index, buffer, shared),
        name="listen_thread"
    ).start()

    print('Starting predition thread...')
    threading.Thread(
        target=lambda: predict(buffer, shared),
        name="predict_thread"
    ).start()

'''
------------------------------------------------------------------------------------------------------------------------------

EXECUTION CODE

------------------------------------------------------------------------------------------------------------------------------
'''

# main method
if __name__ == '__main__':
    print('Initializing Tensorflow Lite models...')
    model_a, input_details, output_details  = init_tflite_model(resource_path('model_a.tflite'))
    model_b, _, _ = init_tflite_model(resource_path('model_b.tflite'))

    print('Starting http server thread...')
    threading.Thread(
        target=lambda: run_http_server(),
        name="snake_http_server"
    ).start()

    print('Starting inference thread...')
    threading.Thread(
        target=lambda: run_inference(),
        name="inference_thread"
    ).start()

    sleep(1)
    webbrowser.open_new("http://localhost:23336/app")


from __future__ import print_function
import os
import time
import tensorflow as tf
import numpy as np
import random
from utils import *
from LIP_model import *

# Set gpus
gpus = [0] # Here I set CUDA to only see one GPU
os.environ["CUDA_VISIBLE_DEVICES"]=','.join([str(i) for i in gpus])
num_gpus = len(gpus) # number of GPUs to use

### parameters setting
N_CLASSES = 20
INPUT_SIZE = (384, 384)
BATCH_SIZE = 1
BATCH_I = 1
SHUFFLE = True
RANDOM_SCALE = True
RANDOM_MIRROR = True
LEARNING_RATE = 1e-4
MOMENTUM = 0.9
POWER = 0.9
NUM_STEPS = 7616 * 35 + 1
SAVE_PRED_EVERY = 7616
p_Weight = 1
s_Weight = 1
DATA_DIR = './datasets/lip'
LIST_PATH = './datasets/lip/list/train_rev.txt'
DATA_ID_LIST = './datasets/lip/list/train_id.txt'
SNAPSHOT_DIR = './checkpoint/JPPNet-s2'
LOG_DIR = './logs/JPPNet-s2'

def main():
    RANDOM_SEED = random.randint(1000, 9999)
    tf.set_random_seed(RANDOM_SEED)

    # Create queue coordinator.
    coord = tf.train.Coordinator()
    h, w = INPUT_SIZE
    # Load reader.
    with tf.name_scope("create_inputs"):
        reader = LIPReader(DATA_DIR, LIST_PATH, DATA_ID_LIST, INPUT_SIZE, RANDOM_SCALE, RANDOM_MIRROR, SHUFFLE, coord)
        image_batch, label_batch, heatmap_batch = reader.dequeue(BATCH_SIZE)
        image_batch075 = tf.image.resize_images(image_batch, [int(h * 0.75), int(w * 0.75)])
        image_batch050 = tf.image.resize_images(image_batch, [int(h * 0.5), int(w * 0.5)])
        heatmap_batch = tf.scalar_mul(1.0/255, heatmap_batch)

    tower_grads = []
    reuse1 = False
    reuse2 = False
    # Define loss and optimisation parameters.
    base_lr = tf.constant(LEARNING_RATE)
    step_ph = tf.placeholder(dtype=tf.float32, shape=())
    learning_rate = tf.scalar_mul(base_lr, tf.pow((1 - step_ph / NUM_STEPS), POWER))
    optim = tf.train.MomentumOptimizer(learning_rate, MOMENTUM)

    for i in xrange (num_gpus):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('Tower_%d' % (i)) as scope:
                if i == 0:
                    reuse1 = False
                    reuse2 = True
                else:
                    reuse1 = True
                    reuse2 = True
                next_image = image_batch[i*BATCH_I:(i+1)*BATCH_I,:]
                next_image075 = image_batch075[i*BATCH_I:(i+1)*BATCH_I,:]
                next_image050 = image_batch050[i*BATCH_I:(i+1)*BATCH_I,:]
                next_heatmap = heatmap_batch[i*BATCH_I:(i+1)*BATCH_I,:]
                next_label = label_batch[i*BATCH_I:(i+1)*BATCH_I,:]

                # Create network.
                with tf.variable_scope('', reuse=reuse1):
                    net_100 = JPPNetModel({'data': next_image}, is_training=False, n_classes=N_CLASSES)
                with tf.variable_scope('', reuse=reuse2):
                    net_075 = JPPNetModel({'data': next_image075}, is_training=False, n_classes=N_CLASSES)
                with tf.variable_scope('', reuse=reuse2):
                    net_050 = JPPNetModel({'data': next_image050}, is_training=False, n_classes=N_CLASSES)

                # parsing net
                parsing_fea1_100 = net_100.layers['res5d_branch2b_parsing']
                parsing_fea1_075 = net_075.layers['res5d_branch2b_parsing']
                parsing_fea1_050 = net_050.layers['res5d_branch2b_parsing']

                parsing_out1_100 = net_100.layers['fc1_human']
                parsing_out1_075 = net_075.layers['fc1_human']
                parsing_out1_050 = net_050.layers['fc1_human']
                # pose net
                resnet_fea_100 = net_100.layers['res4b22_relu']
                resnet_fea_075 = net_075.layers['res4b22_relu']
                resnet_fea_050 = net_050.layers['res4b22_relu']

                with tf.variable_scope('', reuse=reuse1):
                    pose_out1_100, pose_fea1_100 = pose_net(resnet_fea_100, 'fc1_pose')
                    pose_out2_100, pose_fea2_100 = pose_refine(pose_out1_100, parsing_out1_100, pose_fea1_100, name='fc2_pose')
                    parsing_out2_100, parsing_fea2_100 = parsing_refine(parsing_out1_100, pose_out1_100, parsing_fea1_100, name='fc2_parsing')
                    parsing_out3_100, parsing_fea3_100 = parsing_refine(parsing_out2_100, pose_out2_100, parsing_fea2_100, name='fc3_parsing')
                    pose_out3_100, pose_fea3_100 = pose_refine(pose_out2_100, parsing_out2_100, pose_fea2_100, name='fc3_pose')

                with tf.variable_scope('', reuse=reuse2):
                    pose_out1_075, pose_fea1_075 = pose_net(resnet_fea_075, 'fc1_pose')
                    pose_out2_075, pose_fea2_075 = pose_refine(pose_out1_075, parsing_out1_075, pose_fea1_075, name='fc2_pose')
                    parsing_out2_075, parsing_fea2_075 = parsing_refine(parsing_out1_075, pose_out1_075, parsing_fea1_075, name='fc2_parsing')
                    parsing_out3_075, parsing_fea3_075 = parsing_refine(parsing_out2_075, pose_out2_075, parsing_fea2_075, name='fc3_parsing')
                    pose_out3_075, pose_fea3_075 = pose_refine(pose_out2_075, parsing_out2_075, pose_fea2_075, name='fc3_pose')

                with tf.variable_scope('', reuse=reuse2):
                    pose_out1_050, pose_fea1_050 = pose_net(resnet_fea_050, 'fc1_pose')
                    pose_out2_050, pose_fea2_050 = pose_refine(pose_out1_050, parsing_out1_050, pose_fea1_050, name='fc2_pose')
                    parsing_out2_050, parsing_fea2_050 = parsing_refine(parsing_out1_050, pose_out1_050, parsing_fea1_050, name='fc2_parsing')
                    parsing_out3_050, parsing_fea3_050 = parsing_refine(parsing_out2_050, pose_out2_050, parsing_fea2_050, name='fc3_parsing')
                    pose_out3_050, pose_fea3_050 = pose_refine(pose_out2_050, parsing_out2_050, pose_fea2_050, name='fc3_pose')

                # combine resize
                parsing_out1 = tf.reduce_mean(tf.stack([parsing_out1_100,
                                                     tf.image.resize_images(parsing_out1_075, tf.shape(parsing_out1_100)[1:3,]),
                                                     tf.image.resize_images(parsing_out1_050, tf.shape(parsing_out1_100)[1:3,])]), axis=0)
                parsing_out2 = tf.reduce_mean(tf.stack([parsing_out2_100,
                                                     tf.image.resize_images(parsing_out2_075, tf.shape(parsing_out2_100)[1:3,]),
                                                     tf.image.resize_images(parsing_out2_050, tf.shape(parsing_out2_100)[1:3,])]), axis=0)
                parsing_out3 = tf.reduce_mean(tf.stack([parsing_out3_100,
                                                     tf.image.resize_images(parsing_out3_075, tf.shape(parsing_out3_100)[1:3,]),
                                                     tf.image.resize_images(parsing_out3_050, tf.shape(parsing_out3_100)[1:3,])]), axis=0)
                pose_out1 = tf.reduce_mean(tf.stack([pose_out1_100,
                                                     tf.image.resize_nearest_neighbor(pose_out1_075, tf.shape(pose_out1_100)[1:3,]),
                                                     tf.image.resize_nearest_neighbor(pose_out1_050, tf.shape(pose_out1_100)[1:3,])]), axis=0)
                pose_out2 = tf.reduce_mean(tf.stack([pose_out2_100,
                                                     tf.image.resize_nearest_neighbor(pose_out2_075, tf.shape(pose_out2_100)[1:3,]),
                                                     tf.image.resize_nearest_neighbor(pose_out2_050, tf.shape(pose_out2_100)[1:3,])]), axis=0)
                pose_out3 = tf.reduce_mean(tf.stack([pose_out3_100,
                                                     tf.image.resize_nearest_neighbor(pose_out3_075, tf.shape(pose_out3_100)[1:3,]),
                                                     tf.image.resize_nearest_neighbor(pose_out3_050, tf.shape(pose_out3_100)[1:3,])]), axis=0)

                ### Predictions: ignoring all predictions with labels greater or equal than n_classes
                raw_prediction_p1 = tf.reshape(parsing_out1, [-1, N_CLASSES])
                raw_prediction_p1_100 = tf.reshape(parsing_out1_100, [-1, N_CLASSES])
                raw_prediction_p1_075 = tf.reshape(parsing_out1_075, [-1, N_CLASSES])
                raw_prediction_p1_050 = tf.reshape(parsing_out1_050, [-1, N_CLASSES])

                raw_prediction_p2 = tf.reshape(parsing_out2, [-1, N_CLASSES])
                raw_prediction_p2_100 = tf.reshape(parsing_out2_100, [-1, N_CLASSES])
                raw_prediction_p2_075 = tf.reshape(parsing_out2_075, [-1, N_CLASSES])
                raw_prediction_p2_050 = tf.reshape(parsing_out2_050, [-1, N_CLASSES])

                raw_prediction_p3 = tf.reshape(parsing_out3, [-1, N_CLASSES])
                raw_prediction_p3_100 = tf.reshape(parsing_out3_100, [-1, N_CLASSES])
                raw_prediction_p3_075 = tf.reshape(parsing_out3_075, [-1, N_CLASSES])
                raw_prediction_p3_050 = tf.reshape(parsing_out3_050, [-1, N_CLASSES])

                label_proc = prepare_label(next_label, tf.stack(parsing_out1.get_shape()[1:3]), one_hot=False) # [batch_size, h, w]
                label_proc075 = prepare_label(next_label, tf.stack(parsing_out1_075.get_shape()[1:3]), one_hot=False)
                label_proc050 = prepare_label(next_label, tf.stack(parsing_out1_050.get_shape()[1:3]), one_hot=False)

                raw_gt = tf.reshape(label_proc, [-1,])
                raw_gt075 = tf.reshape(label_proc075, [-1,])
                raw_gt050 = tf.reshape(label_proc050, [-1,])

                indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, N_CLASSES - 1)), 1)
                indices075 = tf.squeeze(tf.where(tf.less_equal(raw_gt075, N_CLASSES - 1)), 1)
                indices050 = tf.squeeze(tf.where(tf.less_equal(raw_gt050, N_CLASSES - 1)), 1)

                gt = tf.cast(tf.gather(raw_gt, indices), tf.int32)
                gt075 = tf.cast(tf.gather(raw_gt075, indices075), tf.int32)
                gt050 = tf.cast(tf.gather(raw_gt050, indices050), tf.int32)

                prediction_p1 = tf.gather(raw_prediction_p1, indices)
                prediction_p1_100 = tf.gather(raw_prediction_p1_100, indices)
                prediction_p1_075 = tf.gather(raw_prediction_p1_075, indices075)
                prediction_p1_050 = tf.gather(raw_prediction_p1_050, indices050)

                prediction_p2 = tf.gather(raw_prediction_p2, indices)
                prediction_p2_100 = tf.gather(raw_prediction_p2_100, indices)
                prediction_p2_075 = tf.gather(raw_prediction_p2_075, indices075)
                prediction_p2_050 = tf.gather(raw_prediction_p2_050, indices050)

                prediction_p3 = tf.gather(raw_prediction_p3, indices)
                prediction_p3_100 = tf.gather(raw_prediction_p3_100, indices)
                prediction_p3_075 = tf.gather(raw_prediction_p3_075, indices075)
                prediction_p3_050 = tf.gather(raw_prediction_p3_050, indices050)

                next_heatmap075 = tf.image.resize_nearest_neighbor(next_heatmap, pose_out1_075.get_shape()[1:3])
                next_heatmap050 = tf.image.resize_nearest_neighbor(next_heatmap, pose_out1_050.get_shape()[1:3])

                ### Pixel-wise softmax loss.
                loss_p1 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p1, labels=gt))
                loss_p1_100 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p1_100, labels=gt))
                loss_p1_075 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p1_075, labels=gt075))
                loss_p1_050 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p1_050, labels=gt050))

                loss_p2 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p2, labels=gt))
                loss_p2_100 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p2_100, labels=gt))
                loss_p2_075 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p2_075, labels=gt075))
                loss_p2_050 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p2_050, labels=gt050))

                loss_p3 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p3, labels=gt))
                loss_p3_100 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p3_100, labels=gt))
                loss_p3_075 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p3_075, labels=gt075))
                loss_p3_050 = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction_p3_050, labels=gt050))

                loss_s1 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap, pose_out1)), [1, 2, 3])))
                loss_s1_100 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap, pose_out1_100)), [1, 2, 3])))
                loss_s1_075 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap075, pose_out1_075)), [1, 2, 3])))
                loss_s1_050 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap050, pose_out1_050)), [1, 2, 3])))

                loss_s2 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap, pose_out2)), [1, 2, 3])))
                loss_s2_100 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap, pose_out2_100)), [1, 2, 3])))
                loss_s2_075 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap075, pose_out2_075)), [1, 2, 3])))
                loss_s2_050 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap050, pose_out2_050)), [1, 2, 3])))

                loss_s3 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap, pose_out3)), [1, 2, 3])))
                loss_s3_100 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap, pose_out3_100)), [1, 2, 3])))
                loss_s3_075 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap075, pose_out3_075)), [1, 2, 3])))
                loss_s3_050 = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(next_heatmap050, pose_out3_050)), [1, 2, 3])))

                loss_parsing = loss_p1 + loss_p1_100 + loss_p1_075 + loss_p1_050 + loss_p2 + loss_p2_100 + loss_p2_075 + loss_p2_050 + loss_p3 + loss_p3_100 + loss_p3_075 + loss_p3_050
                loss_pose = loss_s1 + loss_s1_100 + loss_s1_075 + loss_s1_050 + loss_s2 + loss_s2_100 + loss_s2_075 + loss_s2_050 + loss_s3 + loss_s3_100 + loss_s3_075 + loss_s3_050
                reduced_loss =  loss_pose * s_Weight + loss_parsing * p_Weight

                trainable_variable = tf.trainable_variables()
                grads = optim.compute_gradients(reduced_loss, var_list=trainable_variable)

                tower_grads.append(grads)

                tf.add_to_collection('loss_p1', loss_p1)
                tf.add_to_collection('loss_p2', loss_p2)
                tf.add_to_collection('loss_p3', loss_p3)
                tf.add_to_collection('loss_s1', loss_s1)
                tf.add_to_collection('loss_s2', loss_s2)
                tf.add_to_collection('loss_s3', loss_s3)
                tf.add_to_collection('reduced_loss', reduced_loss)

    # Average the gradients
    grads_ave = average_gradients(tower_grads)
    # apply the gradients with our optimizers
    train_op = optim.apply_gradients(grads_ave)

    loss_p1_ave = tf.reduce_mean(tf.get_collection('loss_p1'))
    loss_p2_ave = tf.reduce_mean(tf.get_collection('loss_p2'))
    loss_p3_ave = tf.reduce_mean(tf.get_collection('loss_p3'))
    loss_s1_ave = tf.reduce_mean(tf.get_collection('loss_s1'))
    loss_s2_ave = tf.reduce_mean(tf.get_collection('loss_s2'))
    loss_s3_ave = tf.reduce_mean(tf.get_collection('loss_s3'))
    loss_ave = tf.reduce_mean(tf.get_collection('reduced_loss'))

    loss_summary_p1 = tf.summary.scalar("loss_p1_ave", loss_p1_ave)
    loss_summary_p2 = tf.summary.scalar("loss_p2_ave", loss_p2_ave)
    loss_summary_p3 = tf.summary.scalar("loss_p3_ave", loss_p3_ave)
    loss_summary_s1 = tf.summary.scalar("loss_s1_ave", loss_s1_ave)
    loss_summary_s2 = tf.summary.scalar("loss_s2_ave", loss_s2_ave)
    loss_summary_s3 = tf.summary.scalar("loss_s3_ave", loss_s3_ave)
    loss_summary_ave = tf.summary.scalar("loss_ave", loss_ave)
    loss_summary = tf.summary.merge([loss_summary_ave, loss_summary_s1, loss_summary_s2, loss_summary_s3, loss_summary_p1, loss_summary_p2, loss_summary_p3])
    summary_writer = tf.summary.FileWriter(LOG_DIR, graph=tf.get_default_graph())

    # Set up tf session and initialize variables.
    config = tf.ConfigProto(allow_soft_placement=True,log_device_placement=False)
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()
    sess.run(init)

    # Saver for storing checkpoints of the model.
    all_saver_var = tf.global_variables()
    restore_var = all_saver_var #[v for v in all_saver_var if 'pose' not in v.name and 'parsing' not in v.name]
    saver = tf.train.Saver(var_list=all_saver_var, max_to_keep=50)
    loader = tf.train.Saver(var_list=restore_var)

    if load(loader, sess, SNAPSHOT_DIR):
        print(" [*] Load SUCCESS")
    else:
        print(" [!] Load failed...")

    # Start queue threads.
    threads = tf.train.start_queue_runners(coord=coord, sess=sess)

    # Iterate over training steps.
    for step in range(NUM_STEPS):
        start_time = time.time()
        loss_value = 0
        feed_dict = { step_ph : step }

        # Apply gradients.
        summary, loss_value, _ = sess.run([loss_summary, reduced_loss, train_op], feed_dict=feed_dict)
        summary_writer.add_summary(summary, step)
        if step % SAVE_PRED_EVERY == 0:
            save(saver, sess, SNAPSHOT_DIR, step)

        duration = time.time() - start_time
        print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration))
    coord.request_stop()
    coord.join(threads)

def average_gradients(tower_grads):
  """Calculate the average gradient for each shared variable across all towers.
  Note that this function provides a synchronization point across all towers.
  Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
  Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
  """
  average_grads = []
  for grad_and_vars in zip(*tower_grads):
    # Note that each grad_and_vars looks like the following:
    #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
    grads = []
    for g, _ in grad_and_vars:
      # Add 0 dimension to the gradients to represent the tower.
      expanded_g = tf.expand_dims(g, 0)

      # Append on a 'tower' dimension which we will average over below.
      grads.append(expanded_g)

    # Average over the 'tower' dimension.
    grad = tf.concat(axis=0, values=grads)
    grad = tf.reduce_mean(grad, 0)

    # Keep in mind that the Variables are redundant because they are shared
    # across towers. So .. we will just return the first tower's pointer to
    # the Variable.
    v = grad_and_vars[0][1]
    grad_and_var = (grad, v)
    average_grads.append(grad_and_var)
  return average_grads

if __name__ == '__main__':
    main()


import numpy as np
import torch
from torchvision import datasets, transforms
import torch.optim as optim
import torch.nn.functional as F
from model import Net
from utils import *
import os

LR = [0.01, 0.1, 1.0]

def train(model, device, train_loader, optimizer, epoch, lr, p):

    model.train()

    fs = 44100
    duration = 0.01
    f = 200.0

    frames = []

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()

        norms = []
        for layer in model.ordered_layers:
            norm_grad = layer.weight.grad.norm()
            norms.append(norm_grad)

            tone = f + ((norm_grad.numpy()) * 100.0)
            tone = tone.astype(np.float32)
            samples = generate_tone(fs, tone, duration)

            frames.append(samples)

        silence = np.zeros(samples.shape[0] * 2,
                           dtype=np.float32)
        frames.append(silence)

        optimizer.step()

        # Just 200 steps per epoach
        if batch_idx == 200:
            break

    wf = wave.open("training_audio/sgd_" + str(lr) + ".wav", 'wb')
    wf.setnchannels(1)
    wf.setsampwidth(p.get_sample_size(pyaudio.paFloat32))
    wf.setframerate(fs)
    wf.writeframes(b''.join(frames))
    wf.close()
    print("Saved audio file for SGD LR:", lr)

def main():
    device = torch.device("cpu")

    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=256, shuffle=True)

    model = Net().to(device)

    fs = 44100
    duration = 0.01
    f = 200.0
    p, stream = open_stream(fs)

    if os.path.isdir("training_audio") == False:
        os.mkdir("./training_audio")

    for lr in LR:
        print("Training NN with LR:", lr)
        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.5)

        for epoch in range(1, 2):
            train(model, device, train_loader, optimizer, epoch, lr, p)

    stream.stop_stream()
    stream.close()
    p.terminate()

if __name__ == "__main__":
    main()