import argparse
import collections
import logging
import os
import threading
import time
import timeit
import traceback

import gym
import nest
import torch
import libtorchbeast
from torch import nn
from torch.nn import functional as F
from torchbeast.core import file_writer
from torchbeast.core import vtrace

from torchbeast_procgen.model import Net, DeeperNet
from torchbeast_procgen.model import DeltaNetModel as DeltaNet
from torchbeast_procgen.model import DeltaDeltaNetModel as DDNet
from torchbeast_procgen.model import SRModel as SRM
from torchbeast_procgen.model import PseudoSRModel as PseudoSRM
from torchbeast_procgen.model import NoCarryOverSRModel as NoCarrySRM
from torchbeast_procgen.model import SMFWPModel as SMFWP

from torchbeast_procgen.model import LinearTransformerModel as LT
from torchbeast_procgen.model import RecDeltaModel as RecDelta
from torchbeast_procgen.model import FastRNNModel as FastRNN

# Necessary for multithreading.
os.environ["OMP_NUM_THREADS"] = "1"

# Make sure these are consistent with the lists in polybeast_env
# lexicographically sorted list
list_procgen_env_lex = [
    'bigfish',
    'bossfight',
    'caveflyer',
    'chaser',
    'climber',
    'coinrun',
    'dodgeball',
    'fruitbot',
    'heist',
    'jumper',
    'leaper',
    'maze',
    'miner',
    'ninja',
    'plunder',
    'starpilot',
]

# interesting ones first, we use this list
list_procgen_env = [
    'bigfish',
    'fruitbot',
    'maze',
    'leaper',
    'plunder',
    'starpilot',
    'miner',
    'bossfight',
    'caveflyer',
    'chaser',
    'climber',
    'coinrun',
    'dodgeball',
    'heist',
    'jumper',
    'ninja',
]

# env with memory mode extension
list_procgen_env_mem = [
    'dodgeball',
    'heist',
    'maze',
    'miner',
    'caveflyer',
    'jumper',
]

# yapf: disable
parser = argparse.ArgumentParser(description="PyTorch Scalable Agent")

parser.add_argument("--pipes_basename", default="unix:/tmp/polybeast",
                    help="Basename for the pipes for inter-process communication. "
                    "Has to be of the type unix:/some/path.")
parser.add_argument("--mode", default="train",
                    choices=["train", "test", "test_render"],
                    help="Training or test mode.")
parser.add_argument("--xpid", default=None,
                    help="Experiment id (default: None).")

# Training settings.
parser.add_argument("--single_gpu", action="store_true",
                    help="use single gpu.")
parser.add_argument("--disable_checkpoint", action="store_true",
                    help="Disable saving checkpoint.")
parser.add_argument("--disable_validation", action="store_true",
                    help="Disable validation.")
parser.add_argument("--validate_every", default=10, type=int,
                    help="run validation every this *minutes*.")
parser.add_argument("--validate_step_every", default=-1, type=int,
                    help="run validation every this *steps*.")
parser.add_argument("--save_extra_checkpoint", default=50000000, type=int,
                    help="Save an extra checkpoint at .")
parser.add_argument("--eval_extra", action="store_true",
                    help="Eval extra checkpoint.")
parser.add_argument("--savedir", default="~/palaas/torchbeast",
                    help="Root dir where experiment data will be saved.")
parser.add_argument("--num_actors", default=4, type=int, metavar="N",
                    help="Number of actors.")
parser.add_argument("--total_steps", default=100000, type=int, metavar="T",
                    help="Total environment steps to train for.")
parser.add_argument("--batch_size", default=8, type=int, metavar="B",
                    help="Learner batch size.")
parser.add_argument("--unroll_length", default=80, type=int, metavar="T",
                    help="The unroll length (time dimension).")
parser.add_argument("--num_learner_threads", default=2, type=int,
                    metavar="N", help="Number learner threads.")
parser.add_argument("--num_inference_threads", default=2, type=int,
                    metavar="N", help="Number learner threads.")
parser.add_argument("--disable_cuda", action="store_true",
                    help="Disable CUDA.")
parser.add_argument("--num_actions", default=6, type=int, metavar="A",
                    help="Number of actions.")
parser.add_argument("--conv_scale", default=1, type=int,
                    help="[ff] scale for num channels in conv layers.")
parser.add_argument("--use_lstm", action="store_true",
                    help="Use LSTM in agent model.")
parser.add_argument("--use_deep_ff", action="store_true",
                    help="Use deep FF net in agent model.")
parser.add_argument("--use_delta_rnn", action="store_true",
                    help="Use Delta RNN in agent model.")
parser.add_argument("--use_delta", action="store_true",
                    help="Use Delta Net in agent model.")
parser.add_argument("--use_lt", action="store_true",
                    help="Use Linear Trafo in agent model.")
parser.add_argument("--use_rec_delta", action="store_true",
                    help="Use Recurrent Delta Net in agent model.")
parser.add_argument("--use_dd", action="store_true",
                    help="Use Delta Delta in agent model.")
parser.add_argument("--use_sr", action="store_true",
                    help="Use SR matrix in agent model.")
parser.add_argument("--use_psr", action="store_true",
                    help="Use pseudoSR matrix in agent model.")
parser.add_argument("--use_smfwp", action="store_true",
                    help="Use SR matrix in agent model.")
parser.add_argument("--use_no_carry_sr", action="store_true",
                    help="Use SR matrix w/o carry over in agent model.")
parser.add_argument("--test_no_carry_sr", action="store_true",
                    help="Use SR matrix w/o carry over in test.")
parser.add_argument("--keep_dd", action="store_true",
                    help="Keep delta delta XEM as part of the model.")
parser.add_argument("--max_learner_queue_size", default=None, type=int, metavar="N",
                    help="Optional maximum learner queue size. Defaults to batch_size.")

# Model settings.
parser.add_argument("--hidden_size", default=128, type=int,
                    help="transformer hidden size.")
parser.add_argument("--dim_ff", default=512, type=int,
                    help="transformer hidden size.")
parser.add_argument("--dim_head", default=32, type=int,
                    help="transformer head size.")
parser.add_argument("--num_layers", default=2, type=int,
                    help="tranformer num layers.")
parser.add_argument("--num_head", default=4, type=int,
                    help="tranformer num heads.")
parser.add_argument("--dropout", default=0.0, type=float,
                    help="tranformer dropout.")

# Loss settings.
parser.add_argument("--entropy_cost", default=0.0006, type=float,
                    help="Entropy cost/multiplier.")
parser.add_argument("--baseline_cost", default=0.5, type=float,
                    help="Baseline cost/multiplier.")
parser.add_argument("--discounting", default=0.99, type=float,
                    help="Discounting factor.")
parser.add_argument("--reward_clipping", default="abs_one",
                    choices=["abs_one", "none"],
                    help="Reward clipping.")

# Optimizer settings.
parser.add_argument("--learning_rate", default=0.00048, type=float,
                    metavar="LR", help="Learning rate.")
parser.add_argument("--alpha", default=0.99, type=float,
                    help="RMSProp smoothing constant.")
parser.add_argument("--momentum", default=0, type=float,
                    help="RMSProp momentum.")
parser.add_argument("--epsilon", default=0.01, type=float,
                    help="RMSProp epsilon.")
parser.add_argument("--grad_norm_clipping", default=40.0, type=float,
                    help="Global gradient norm clip.")

# Misc settings.
parser.add_argument("--write_profiler_trace", action="store_true",
                    help="Collect and write a profiler trace "
                    "for chrome://tracing/.")
# yapf: enable

parser.add_argument('--num_servers', default=4, type=int, metavar='N',
                    help='Number of environment servers.')
parser.add_argument('--env', type=str, default='PongNoFrameskip-v4',
                    help='Gym environment.')
parser.add_argument('--multi_env', default=1, type=int, metavar='N',
                    help='number of env to jointly train on.')
parser.add_argument('--allow_oov', action="store_true",
                    help='Allow action space larger than the env specific one.'
                    ' All out-of-vocab action will be mapped to NoOp.')
parser.add_argument('--num_levels', default=0, type=int, metavar='N',
                    help='Procgen num_levels.')
parser.add_argument('--start_level', default=0, type=int, metavar='N',
                    help='Procgen start_level.')
parser.add_argument('--distribution_mode', type=str, default='hard',
                    choices=[
                        'easy', 'hard', 'extreme', 'memory', 'exploration'],
                    help='distribution mode.')
parser.add_argument('--valid_num_levels', default=0, type=int, metavar='N',
                    help='Procgen num_levels for validation set.')
parser.add_argument('--valid_start_level', default=0, type=int, metavar='N',
                    help='Procgen start_level for validation set.')
parser.add_argument('--valid_num_episodes', default=5, type=int, metavar='N',
                    help='number of validation episodes.')
parser.add_argument('--valid_num_runs', default=1, type=int, metavar='N',
                    help='number of validation runs '
                         '(each run is valid_num_episodes episode long).')
parser.add_argument('--valid_distribution_mode', type=str, default='hard',
                    choices=[
                        'easy', 'hard', 'extreme', 'memory', 'exploration'],
                    help='validation distribution mode.')

# For eval
parser.add_argument('--test_model_name', type=str,
                    help='model checkpoint suffix for evaluation.')
parser.add_argument('--test_num_levels', default=0, type=int, metavar='N',
                    help='Procgen num_levels for validation set.')
parser.add_argument('--test_start_level', default=0, type=int, metavar='N',
                    help='Procgen start_level for validation set.')
parser.add_argument('--test_distribution_mode', type=str, default='hard',
                    choices=[
                        'easy', 'hard', 'extreme', 'memory', 'exploration'],
                    help='test distribution mode.')

# Wandb settings
parser.add_argument('--project_name', type=str, default=None,
                    help='project name for wandb.')
parser.add_argument('--job_name', type=str, default=None,
                    help='job name for wandb.')
parser.add_argument('--use_wandb', action='store_true',
                    help='use wandb.')

args = parser.parse_args()

if args.use_wandb:  # configure wandb.
    import wandb
    use_wandb = True

    if args.project_name is None:
        project_name = (os.uname()[1]
                        + datetime.now().strftime("%Y-%m-%d-%H-%M-%S"))
    else:
        project_name = args.project_name

    wandb.init(project=project_name)

    if args.job_name is None:
        # wandb.run.name = (os.uname()[1]
        #                   + datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        #                   + args.work_dir)
        wandb.run.name = f"{os.uname()[1]}" \
                         f"-{args.mode}" \
                         f"-{args.xpid}" \
                         f"-{args.disable_checkpoint}" \
                         f"-{args.savedir}" \
                         f"-{args.num_actors}" \
                         f"-{args.total_steps}" \
                         f"-{args.batch_size}" \
                         f"-{args.unroll_length}" \
                         f"-{args.entropy_cost}" \
                         f"-{args.baseline_cost}" \
                         f"-{args.discounting}" \
                         f"-{args.reward_clipping}" \
                         f"-{args.learning_rate}" \
                         f"-{args.alpha}" \
                         f"-{args.momentum}" \
                         f"-{args.epsilon}" \
                         f"-{args.grad_norm_clipping}"
    else:
        wandb.run.name = f"{os.uname()[1]}//{args.job_name}"

    config = wandb.config
    config.host = os.uname()[1]  # host node name
    config.mode=args.mode
    config.xpid=args.xpid
    config.disable_checkpoint=args.disable_checkpoint
    config.savedir=args.savedir
    config.num_actors=args.num_actors
    config.total_steps=args.total_steps
    config.batch_size=args.batch_size
    config.unroll_length=args.unroll_length
    config.disable_cuda=args.disable_cuda
    config.use_lstm=args.use_lstm
    config.entropy_cost=args.entropy_cost
    config.baseline_cost=args.baseline_cost
    config.discounting=args.discounting
    config.reward_clipping=args.reward_clipping
    config.learning_rate=args.learning_rate
    config.alpha=args.alpha
    config.momentum=args.momentum
    config.epsilon=args.epsilon
    config.grad_norm_clipping=args.grad_norm_clipping
else:
    use_wandb = False

logging.basicConfig(
    format=(
        "[%(levelname)s:%(process)d %(module)s:%(lineno)d %(asctime)s] " "%(message)s"
    ),
    level=0,
)

def compute_baseline_loss(advantages):
    return 0.5 * torch.sum(advantages ** 2)

def compute_entropy_loss(logits):
    """Return the entropy loss, i.e., the negative entropy of the policy."""
    policy = F.softmax(logits, dim=-1)
    log_policy = F.log_softmax(logits, dim=-1)
    return torch.sum(policy * log_policy)

def compute_policy_gradient_loss(logits, actions, advantages):
    cross_entropy = F.nll_loss(
        F.log_softmax(torch.flatten(logits, 0, 1), dim=-1),
        target=torch.flatten(actions, 0, 1),
        reduction="none",
    )
    cross_entropy = cross_entropy.view_as(advantages)
    return torch.sum(cross_entropy * advantages.detach())

def inference(flags, inference_batcher, model, lock=threading.Lock()):  # noqa: B008
    with torch.no_grad():
        for batch in inference_batcher:
            batched_env_outputs, agent_state = batch.get_inputs()
            frame, reward, done, *_ = batched_env_outputs
            frame = frame.to(flags.actor_device, non_blocking=True)
            reward = reward.to(flags.actor_device, non_blocking=True)
            done = done.to(flags.actor_device, non_blocking=True)
            agent_state = nest.map(
                lambda t: t.to(flags.actor_device, non_blocking=True), agent_state
            )
            with lock:
                outputs = model(
                    dict(frame=frame, reward=reward, done=done), agent_state
                )

            outputs = nest.map(lambda t: t.cpu(), outputs)
            batch.set_outputs(outputs)

EnvOutput = collections.namedtuple(
    "EnvOutput", "frame rewards done episode_step episode_return"
)
AgentOutput = collections.namedtuple("AgentOutput", "action policy_logits baseline")
Batch = collections.namedtuple("Batch", "env agent")

def learn(
    flags,
    learner_queue,
    model,
    actor_model,
    optimizer,
    scheduler,
    stats,
    plogger,
    lock=threading.Lock(),
):
    for tensors in learner_queue:
        tensors = nest.map(lambda t: t.to(flags.learner_device), tensors)

        batch, initial_agent_state = tensors
        env_outputs, actor_outputs = batch
        frame, reward, done, *_ = env_outputs

        lock.acquire()  # Only one thread learning at a time.
        learner_outputs, unused_state = model(
            dict(frame=frame, reward=reward, done=done), initial_agent_state
        )

        # Take final value function slice for bootstrapping.
        learner_outputs = AgentOutput._make(learner_outputs)
        bootstrap_value = learner_outputs.baseline[-1]

        # Move from obs[t] -> action[t] to action[t] -> obs[t].
        batch = nest.map(lambda t: t[1:], batch)
        learner_outputs = nest.map(lambda t: t[:-1], learner_outputs)

        # Turn into namedtuples again.
        env_outputs, actor_outputs = batch
        env_outputs = EnvOutput._make(env_outputs)
        actor_outputs = AgentOutput._make(actor_outputs)
        learner_outputs = AgentOutput._make(learner_outputs)

        if flags.reward_clipping == "abs_one":
            clipped_rewards = torch.clamp(env_outputs.rewards, -1, 1)
        elif flags.reward_clipping == "none":
            clipped_rewards = env_outputs.rewards

        discounts = (~env_outputs.done).float() * flags.discounting

        vtrace_returns = vtrace.from_logits(
            behavior_policy_logits=actor_outputs.policy_logits,
            target_policy_logits=learner_outputs.policy_logits,
            actions=actor_outputs.action,
            discounts=discounts,
            rewards=clipped_rewards,
            values=learner_outputs.baseline,
            bootstrap_value=bootstrap_value,
        )

        pg_loss = compute_policy_gradient_loss(
            learner_outputs.policy_logits,
            actor_outputs.action,
            vtrace_returns.pg_advantages,
        )
        baseline_loss = flags.baseline_cost * compute_baseline_loss(
            vtrace_returns.vs - learner_outputs.baseline
        )
        entropy_loss = flags.entropy_cost * compute_entropy_loss(
            learner_outputs.policy_logits
        )

        total_loss = pg_loss + baseline_loss + entropy_loss

        optimizer.zero_grad()
        total_loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), flags.grad_norm_clipping)
        optimizer.step()
        scheduler.step()

        actor_model.load_state_dict(model.state_dict())

        episode_returns = env_outputs.episode_return[env_outputs.done]
        stats["step"] = stats.get("step", 0) + flags.unroll_length * flags.batch_size
        stats["episode_returns"] = tuple(episode_returns.cpu().numpy())
        stats["mean_episode_return"] = torch.mean(episode_returns).item()
        stats["mean_episode_step"] = torch.mean(env_outputs.episode_step.float()).item()
        stats["total_loss"] = total_loss.item()
        stats["pg_loss"] = pg_loss.item()
        stats["baseline_loss"] = baseline_loss.item()
        stats["entropy_loss"] = entropy_loss.item()

        stats["learner_queue_size"] = learner_queue.size()

        if use_wandb:
            wandb.log({"episode_returns": stats["episode_returns"]})
            wandb.log({"mean_episode_step": stats["mean_episode_step"]})
            wandb.log({"mean_episode_return": stats["mean_episode_return"]})
            wandb.log({"total_loss": stats["total_loss"]})
            wandb.log({"pg_loss": stats["pg_loss"]})
            wandb.log({"baseline_loss": stats["baseline_loss"]})
            wandb.log({"entropy_loss": stats["entropy_loss"]})

        plogger.log(stats)

        if not len(episode_returns):
            # Hide the mean-of-empty-tuple NaN as it scares people.
            stats["mean_episode_return"] = None

        lock.release()

def train(flags):
    if flags.xpid is None:
        flags.xpid = "torchbeast-%s" % time.strftime("%Y%m%d-%H%M%S")
    plogger = file_writer.FileWriter(
        xpid=flags.xpid, xp_args=flags.__dict__, rootdir=flags.savedir
    )
    checkpointpath = os.path.expandvars(
        os.path.expanduser(
            "%s/%s/%s" % (flags.savedir, flags.xpid, "model.tar"))
    )
    best_valid_checkpointpath = os.path.expandvars(
        os.path.expanduser(
            "%s/%s/%s" % (flags.savedir, flags.xpid, "model_best_val.tar"))
    )

    if flags.save_extra_checkpoint > 0:
        extra_checkpointpath = os.path.expandvars(
            os.path.expanduser(
                "%s/%s/%s" % (flags.savedir, flags.xpid, "model_extra.tar")))

    if flags.single_gpu:
        logging.info("Using single GPU.")
        flags.learner_device = torch.device("cuda:0")
        flags.actor_device = torch.device("cuda:0")
    elif not flags.disable_cuda and torch.cuda.is_available():
        logging.info("Using CUDA.")
        flags.learner_device = torch.device("cuda:0")
        flags.actor_device = torch.device("cuda:1")
    else:
        logging.info("Not using CUDA.")
        flags.learner_device = torch.device("cpu")
        flags.actor_device = torch.device("cpu")

    if flags.max_learner_queue_size is None:
        flags.max_learner_queue_size = flags.batch_size

    # The queue the learner threads will get their data from.
    # Setting `minimum_batch_size == maximum_batch_size`
    # makes the batch size static.
    learner_queue = libtorchbeast.BatchingQueue(
        batch_dim=1,
        minimum_batch_size=flags.batch_size,
        maximum_batch_size=flags.batch_size,
        check_inputs=True,
        maximum_queue_size=flags.max_learner_queue_size,
    )

    # The "batcher", a queue for the inference call. Will yield
    # "batch" objects with `get_inputs` and `set_outputs` methods.
    # The batch size of the tensors will be dynamic.
    inference_batcher = libtorchbeast.DynamicBatcher(
        batch_dim=1,
        minimum_batch_size=1,
        maximum_batch_size=512,
        timeout_ms=100,
        check_outputs=True,
    )

    addresses = []
    connections_per_server = 1
    pipe_id = 0
    while len(addresses) < flags.num_actors:
        for _ in range(connections_per_server):
            addresses.append(f"{flags.pipes_basename}.{pipe_id}")
            if len(addresses) == flags.num_actors:
                break
        pipe_id += 1

    if flags.use_delta:
        model = DeltaNet(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = DeltaNet(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_lt:
        model = LT(num_actions=flags.num_actions, dim_head=flags.dim_head,
                   hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                   num_layers=flags.num_layers, num_head=flags.num_head,
                   dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = LT(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_delta_rnn:
        model = FastRNN(num_actions=flags.num_actions, dim_head=flags.dim_head,
                        hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                        num_layers=flags.num_layers, num_head=flags.num_head,
                        dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = FastRNN(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_rec_delta:
        model = RecDelta(num_actions=flags.num_actions, dim_head=flags.dim_head,
                         hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                         num_layers=flags.num_layers, num_head=flags.num_head,
                         dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = RecDelta(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_dd:
        model = DDNet(num_actions=flags.num_actions, dim_head=flags.dim_head,
                      hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                      num_layers=flags.num_layers, num_head=flags.num_head,
                      dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = DDNet(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_sr:
        model = SRM(num_actions=flags.num_actions, dim_head=flags.dim_head,
                    hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                    num_layers=flags.num_layers, num_head=flags.num_head,
                    dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = SRM(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_psr:
        model = PseudoSRM(num_actions=flags.num_actions, dim_head=flags.dim_head,
                    hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                    num_layers=flags.num_layers, num_head=flags.num_head,
                    dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = PseudoSRM(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_smfwp:
        model = SMFWP(num_actions=flags.num_actions, dim_head=flags.dim_head,
                      hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                      num_layers=flags.num_layers, num_head=flags.num_head,
                      dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = SMFWP(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_no_carry_sr:
        model = NoCarrySRM(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = NoCarrySRM(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    elif flags.use_deep_ff:
        model = DeeperNet(num_actions=flags.num_actions, use_lstm=flags.use_lstm,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, dropout=flags.dropout)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = DeeperNet(num_actions=flags.num_actions, use_lstm=flags.use_lstm,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, dropout=flags.dropout)
        actor_model.to(device=flags.actor_device)

    else:
        model = Net(num_actions=flags.num_actions, conv_scale=flags.conv_scale,
                    use_lstm=flags.use_lstm)
        logging.info(model)
        model = model.to(device=flags.learner_device)

        actor_model = Net(
            num_actions=flags.num_actions, conv_scale=flags.conv_scale,
            use_lstm=flags.use_lstm)
        actor_model.to(device=flags.actor_device)

    # The ActorPool that will run `flags.num_actors` many loops.
    actors = libtorchbeast.ActorPool(
        unroll_length=flags.unroll_length,
        learner_queue=learner_queue,
        inference_batcher=inference_batcher,
        env_server_addresses=addresses,
        initial_agent_state=actor_model.initial_state(),
    )

    def run():
        try:
            actors.run()
        except Exception as e:
            logging.error("Exception in actorpool thread!")
            traceback.print_exc()
            print()
            raise e

    actorpool_thread = threading.Thread(target=run, name="actorpool-thread")

    optimizer = torch.optim.RMSprop(
        model.parameters(),
        lr=flags.learning_rate,
        momentum=flags.momentum,
        eps=flags.epsilon,
        alpha=flags.alpha,
    )

    def lr_lambda(epoch):
        return (
            1
            - min(epoch * flags.unroll_length * flags.batch_size, flags.total_steps)
            / flags.total_steps
        )

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    stats = {}

    best_valid_return = 0

    # Load state from a checkpoint, if possible.
    if os.path.exists(checkpointpath):
        checkpoint_states = torch.load(
            checkpointpath, map_location=flags.learner_device
        )
        model.load_state_dict(checkpoint_states["model_state_dict"])
        optimizer.load_state_dict(checkpoint_states["optimizer_state_dict"])
        scheduler.load_state_dict(checkpoint_states["scheduler_state_dict"])
        stats = checkpoint_states["stats"]
        best_valid_return = checkpoint_states["best_valid_return"]
        logging.info(f"Resuming preempted job, current stats:\n{stats}")

    # Initialize actor model like learner model.
    actor_model.load_state_dict(model.state_dict())

    learner_threads = [
        threading.Thread(
            target=learn,
            name="learner-thread-%i" % i,
            args=(
                flags,
                learner_queue,
                model,
                actor_model,
                optimizer,
                scheduler,
                stats,
                plogger,
            ),
        )
        for i in range(flags.num_learner_threads)
    ]
    inference_threads = [
        threading.Thread(
            target=inference,
            name="inference-thread-%i" % i,
            args=(flags, inference_batcher, actor_model),
        )
        for i in range(flags.num_inference_threads)
    ]

    actorpool_thread.start()
    for t in learner_threads + inference_threads:
        t.start()

    def checkpoint(path):
        if flags.disable_checkpoint:
            return
        logging.info("Saving checkpoint to %s", path)
        torch.save(
            {
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "scheduler_state_dict": scheduler.state_dict(),
                "stats": stats,
                "flags": vars(flags),
                "best_valid_return": best_valid_return,
            },
            path,
        )

    def old_validate(num_episodes=5, num_runs=1):
        from torchbeast_procgen import procgen_wrappers
        from torchbeast.core import environment

        def create_valid_env():
            # better always validate on the same levels via rand_seed?
            return procgen_wrappers.wrap_pytorch(
                procgen_wrappers.wrap_deepmind(
                    gym.make(flags.env,
                             num_levels=flags.valid_num_levels,
                             start_level=flags.valid_start_level,
                             distribution_mode=flags.valid_distribution_mode,
                             rand_seed=None),
                    clip_rewards=False,
                )
            )

        gym_env = create_valid_env()
        env = environment.Environment(gym_env)

        # model.eval()
        device = flags.learner_device  # move to actor device?? which should be less busy? and copy back

        observation = env.initial()
        all_returns = []

        core_state = model.initial_state()
        core_state = nest.map(lambda t: t.to(device), core_state)

        logging.info("Start validation")
        with torch.no_grad():
            while len(all_returns) < num_runs:
                returns = []
                while len(returns) < num_episodes:
                    # (action, policy_logits, baseline), core_state
                    observation = nest.map(lambda t: t.to(device), observation)
                    agent_outputs, core_state = model(observation, core_state)
                    action, _, _ = agent_outputs
                    # action = action.to('cpu')
                    # policy_outputs, _ = agent_outputs
                    observation = env.step(action)
                    # observation = env.step(policy_outputs["action"])
                    if observation["done"].item():
                        returns.append(observation["episode_return"].item())
                        logging.info(
                            "Episode ended after %d steps. Return: %.1f",
                            observation["episode_step"].item(),
                            observation["episode_return"].item(),
                        )
                logging.info(
                    "Average returns over %i episodes: %.1f",
                        num_episodes, sum(returns) / len(returns)
                )
                all_returns.append(sum(returns) / len(returns))
        env.close()
        import numpy as np
        mean_valid_return = np.mean(all_returns)
        logging.info(
            f"[validation] Average returns over "
            f"{num_episodes} for {num_runs} runs: {all_returns}")
        logging.info(
            f"[validation] Mean return: "
            f"{mean_valid_return:.1f}, std: {np.std(all_returns):.1f}")
        model.train()

        if use_wandb:
            wandb.log({"validation return": mean_valid_return})

        return mean_valid_return

    # validate in all env
    def validate(num_episodes=5, num_runs=1):
        from torchbeast_procgen import procgen_wrappers
        from torchbeast.core import environment

        def create_valid_env(env_name):
            # better always validate on the same levels via rand_seed?
            return procgen_wrappers.wrap_pytorch(
                procgen_wrappers.wrap_deepmind(
                    gym.make(env_name,
                             num_levels=flags.valid_num_levels,
                             start_level=flags.valid_start_level,
                             distribution_mode=flags.valid_distribution_mode,
                             rand_seed=None),
                    clip_rewards=False,
                )
            )

        if flags.distribution_mode == 'memory':  # for multi_env training
            list_env = list_procgen_env_mem
        else:
            list_env = list_procgen_env

        cross_env_avg = 0.

        for i in range(flags.multi_env):
            if flags.multi_env == 1:
                env_name = flags.env
                gym_env = create_valid_env(f"{env_name}")
                is_multi_task = False
            else:
                is_multi_task = True
                env_name = list_env[i]
                gym_env = create_valid_env(f"procgen:procgen-{env_name}-v0")
            env = environment.Environment(gym_env)

            # model.eval()
            device = flags.learner_device  # move to actor device?? which should be less busy? and copy back

            observation = env.initial()
            all_returns = []

            core_state = model.initial_state()
            core_state = nest.map(lambda t: t.to(device), core_state)

            logging.info(f"Start validation on {env_name}")
            with torch.no_grad():
                while len(all_returns) < num_runs:
                    returns = []
                    while len(returns) < num_episodes:
                        # (action, policy_logits, baseline), core_state
                        observation = nest.map(
                            lambda t: t.to(device), observation)
                        agent_outputs, core_state = model(
                            observation, core_state)
                        action, _, _ = agent_outputs
                        observation = env.step(action)
                        if observation["done"].item():
                            returns.append(
                                observation["episode_return"].item())
                            logging.info(
                                "Episode ended after %d steps. Return: %.1f",
                                observation["episode_step"].item(),
                                observation["episode_return"].item(),
                            )
                    logging.info(
                        "Average returns over %i episodes: %.1f",
                            num_episodes, sum(returns) / len(returns)
                    )
                    all_returns.append(sum(returns) / len(returns))
            env.close()
            import numpy as np

            if num_runs == 1:
                mean_valid_return = np.mean(returns)
                logging.info(
                    f"[validation, {env_name}] Average returns over "
                    f"{num_episodes} for {num_runs} runs: {returns}")
                logging.info(
                    f"[validation, {env_name}] Mean return: "
                    f"{mean_valid_return:.1f}, std: {np.std(returns):.1f}")
            else:
                mean_valid_return = np.mean(all_returns)
                logging.info(
                    f"[validation, {env_name}] Average returns over "
                    f"{num_episodes} for {num_runs} runs: {all_returns}")
                logging.info(
                    f"[validation, {env_name}] Mean return: "
                    f"{mean_valid_return:.1f}, std: {np.std(all_returns):.1f}")
            if use_wandb:
                wandb.log({f"{env_name} validation return": mean_valid_return})

            cross_env_avg += mean_valid_return
        model.train()

        if is_multi_task:
            cross_env_avg = cross_env_avg / flags.multi_env
            logging.info(
                f"[validation, joint] Cross env avg return {cross_env_avg}")
            if use_wandb:
                wandb.log({"Cross env validation return": cross_env_avg})

        return cross_env_avg

    def format_value(x):
        return f"{x:1.5}" if isinstance(x, float) else str(x)

    try:
        if flags.save_extra_checkpoint > 0:
            saved_extra = False
        last_checkpoint_time = timeit.default_timer()
        while True:
            start_time = timeit.default_timer()
            start_step = stats.get("step", 0)

            if start_step >= flags.total_steps:
                break
            time.sleep(5)
            end_step = stats.get("step", 0)

            if timeit.default_timer() - last_checkpoint_time > flags.validate_every * 60 or start_step % flags.validate_step_every == 0:
                # Validate every 10 min.
                if not flags.disable_validation:
                    val_start = timeit.default_timer()
                    mean_valid_return = validate(
                        num_episodes=flags.valid_num_episodes,
                        num_runs=flags.valid_num_runs)
                    val_dur = val_start - timeit.default_timer()
                    logging.info(f"[validation] duration: "
                                 f"{val_dur:.1f} sec.")
                    if use_wandb:
                        wandb.log({"validation time": val_dur})
                    if mean_valid_return > best_valid_return:
                        # save if best
                        best_valid_return = mean_valid_return
                        checkpoint(best_valid_checkpointpath)
                # Always save latest checkpoint
                checkpoint(checkpointpath)
                last_checkpoint_time = timeit.default_timer()

            logging.info(
                "Step %i @ %.1f SPS. Inference batcher size: %i."
                " Learner queue size: %i."
                " Other stats: (%s)",
                end_step,
                (end_step - start_step) / (timeit.default_timer() - start_time),
                inference_batcher.size(),
                learner_queue.size(),
                ", ".join(
                    f"{key} = {format_value(value)}" for key, value in stats.items()
                ),
            )
            if flags.save_extra_checkpoint > 0:
                if saved_extra is False and end_step > flags.save_extra_checkpoint:
                    logging.info(f"Step {end_step} Saving EXTRA checkpoint to {extra_checkpointpath}")
                    torch.save(
                        {
                            "model_state_dict": model.state_dict(),
                            "optimizer_state_dict": optimizer.state_dict(),
                            "scheduler_state_dict": scheduler.state_dict(),
                            "stats": stats,
                            "flags": vars(flags),
                            "best_valid_return": best_valid_return,
                        },
                        extra_checkpointpath,
                    )
                    saved_extra = True

    except KeyboardInterrupt:
        pass  # Close properly.
    else:
        logging.info("Learning finished after %i steps.", stats["step"])
        checkpoint(checkpointpath)

    # Done with learning. Stop all the ongoing work.
    inference_batcher.close()
    learner_queue.close()

    actorpool_thread.join()

    for t in learner_threads + inference_threads:
        t.join()

# def test(flags, num_episodes=30, num_runs=5, device='cuda'):
def test(flags, num_episodes=200, num_runs=1, device='cuda'):
    if flags.xpid is None:
        checkpointpath = "./latest/model.tar"
    elif flags.eval_extra:
        checkpointpath = os.path.expandvars(
            os.path.expanduser(
                "%s/%s/%s" % (flags.savedir, flags.xpid, "model_extra.tar"))
        )
    elif flags.test_model_name:
        checkpointpath = os.path.expandvars(
            os.path.expanduser(
                "%s/%s/%s" % (
                    flags.savedir, flags.xpid, flags.test_model_name))
        )
    else:
        checkpointpath = os.path.expandvars(
            os.path.expanduser(
                "%s/%s/%s" % (flags.savedir, flags.xpid, "model.tar"))
        )

    from torchbeast.core import environment
    from torchbeast_procgen import procgen_wrappers

    def create_test_env(env_name, num_levels=0, start_level=0,
                        distribution_mode="hard", rand_seed=None,
                        lock=threading.Lock()):
        with lock:  # Atari isn't threadsafe at construction time.
            return procgen_wrappers.wrap_pytorch(
                procgen_wrappers.wrap_deepmind(
                    gym.make(env_name,
                             num_levels=num_levels,
                             start_level=start_level,
                             distribution_mode=distribution_mode,
                             rand_seed=rand_seed),
                    clip_rewards=False,
                )
            )

    # gym_env = create_test_env(flags)
    gym_env = create_test_env(flags.env, num_levels=flags.test_num_levels,
                              start_level=flags.test_start_level,
                              distribution_mode=flags.test_distribution_mode)
    env = environment.Environment(gym_env)

    if flags.use_delta:
        model = DeltaNet(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)

    elif flags.use_lt:
        model = LT(num_actions=flags.num_actions, dim_head=flags.dim_head,
                   hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                   num_layers=flags.num_layers, num_head=flags.num_head,
                   dropout=flags.dropout)

    elif flags.use_delta_rnn:
        model = FastRNN(num_actions=flags.num_actions, dim_head=flags.dim_head,
                        hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                        num_layers=flags.num_layers, num_head=flags.num_head,
                        dropout=flags.dropout)

    elif flags.use_rec_delta:
        model = RecDelta(num_actions=flags.num_actions, dim_head=flags.dim_head,
                         hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                         num_layers=flags.num_layers, num_head=flags.num_head,
                         dropout=flags.dropout)

    elif flags.use_dd:
        model = DDNet(num_actions=flags.num_actions, dim_head=flags.dim_head,
                      hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                      num_layers=flags.num_layers, num_head=flags.num_head,
                      dropout=flags.dropout)

    elif flags.use_sr:
        model = SRM(num_actions=flags.num_actions, dim_head=flags.dim_head,
                    hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                    num_layers=flags.num_layers, num_head=flags.num_head,
                    dropout=flags.dropout)

    elif flags.use_psr:
        model = PseudoSRM(num_actions=flags.num_actions, dim_head=flags.dim_head,
                    hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                    num_layers=flags.num_layers, num_head=flags.num_head,
                    dropout=flags.dropout)

    elif flags.use_no_carry_sr:
        model = NoCarrySRM(
            num_actions=flags.num_actions, dim_head=flags.dim_head,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, num_head=flags.num_head,
            dropout=flags.dropout)

    elif flags.use_smfwp:
        model = SMFWP(num_actions=flags.num_actions, dim_head=flags.dim_head,
                      hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
                      num_layers=flags.num_layers, num_head=flags.num_head,
                      dropout=flags.dropout)

    elif flags.use_deep_ff:
        model = DeeperNet(num_actions=flags.num_actions, use_lstm=flags.use_lstm,
            hidden_size=flags.hidden_size, dim_ff=flags.dim_ff,
            num_layers=flags.num_layers, dropout=flags.dropout)
    else:
        model = Net(num_actions=flags.num_actions, conv_scale=flags.conv_scale,
                    use_lstm=flags.use_lstm)

    print(model)
    print(f"# params: "
          f"{sum(p.numel() for p in model.parameters() if p.requires_grad)}")
    model = model.to(device)
    model.eval()
    checkpoint = torch.load(checkpointpath, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])

    observation = env.initial()
    # returns = []
    all_returns = []

    core_state = model.initial_state()
    core_state = nest.map(lambda t: t.to(device), core_state)

    logging.info("Start eval")
    with torch.no_grad():
        while len(all_returns) < num_runs:
            returns = []
            step_counter = 0
            while len(returns) < num_episodes:
                if flags.mode == "test_render":
                    env.gym_env.render()
                # (action, policy_logits, baseline), core_state
                observation = nest.map(lambda t: t.to(device), observation)
                agent_outputs, core_state = model(observation, core_state)
                step_counter += 1
                if flags.test_no_carry_sr and step_counter == flags.unroll_length:
                    # need to reset state every unroll_length steps
                    step_counter = 0
                    core_state = model.initial_state()
                    core_state = nest.map(lambda t: t.to(device), core_state)
                action, _, _ = agent_outputs
                # action = action.to('cpu')
                # policy_outputs, _ = agent_outputs
                observation = env.step(action)
                # observation = env.step(policy_outputs["action"])
                if observation["done"].item():
                    returns.append(observation["episode_return"].item())
                    logging.info(
                        "Episode ended after %d steps. Return: %.1f",
                        observation["episode_step"].item(),
                        observation["episode_return"].item(),
                    )
            logging.info(
                "Average returns over %i episodes: %.1f",
                    num_episodes, sum(returns) / len(returns)
            )
            all_returns.append(sum(returns) / len(returns))
    env.close()
    import numpy as np
    logging.info(f"Average returns over {num_episodes} for {num_runs} runs: {all_returns}")
    logging.info(f"{flags.env}, Mean return: {np.mean(all_returns):.1f}, std: {np.std(all_returns):.1f}")

def main(flags):
    if not flags.pipes_basename.startswith("unix:"):
        raise Exception("--pipes_basename has to be of the form unix:/some/path.")

    if flags.mode == "train":
        if flags.write_profiler_trace:
            logging.info("Running with profiler.")
            with torch.autograd.profiler.profile() as prof:
                train(flags)
            filename = "chrome-%s.trace" % time.strftime("%Y%m%d-%H%M%S")
            logging.info("Writing profiler trace to '%s.gz'", filename)
            prof.export_chrome_trace(filename)
            os.system("gzip %s" % filename)
        else:
            train(flags)
    else:
        test(flags)

if __name__ == "__main__":
    flags = parser.parse_args()
    main(flags)


import random
import numpy as np
from config import Config
from node import Node
from connection import Connection
from activation import Activation

class Genome() :

    def __init__(self,nodes,connections,id) :
        self.nodes = nodes
        self.connections = connections
        self.fitness = 0
        self.id = id
        self.__divideNodes()

    def __divideNodes(self) :
        self.input_nodes = {}
        self.hidden_nodes = {}
        self.output_nodes = {}
        self.input_connections = {}

        for id,node in self.nodes.items() :
            if(node.type == 'input') :
                self.input_nodes[id] = node
            elif(node.type == 'hidden') :
                self.hidden_nodes[id] = node
            else :
                self.output_nodes[id] = node

        self.input_connections = Genome.__getNodesInputConnections(self.connections)

    @staticmethod
    def __getNodesInputConnections(connections) :
        input_connections = {}
        for _,connection in connections.items() :
            output = connection.output
            if(output.id not in input_connections) :
                input_connections[output.id] = []
            input_connections[output.id].append(connection)
        return input_connections

    def create_model(self) :
        self.diversity_reward = 0
        self.actions = []
        return self

    def run(self,inputs) :
        for _,node in self.nodes.items() :
            node.value = None

        outputs = [0] * len(self.output_nodes)
        for id,node in self.output_nodes.items() :
            node.value = self.__calculateNodeValue(node,inputs)
            if(node.type == 'output') :
                outputs[node.order] = node.value

        return outputs

    def __calculateNodeValue(self,node,inputs) :
        if(node.type == 'input') :
            node.value = inputs[node.order]
            return node.activation(node.value)
        if(node.id not in self.input_connections) :
            node.value = 0
            return node.value
        connections = self.input_connections[node.id]
        value = 0
        for connection in connections :
            if(connection.enabled) :
                input_node = self.nodes[connection.input.id]
                if(input_node.value is None) :
                    input_node.value = self.__calculateNodeValue(input_node,inputs)
                value += input_node.value * connection.weight

        return node.activation(value)

    def distance(self,genome) :
        if(not self.connections) :
            if(genome.connections) :
                return Config.c1 * len(genome.connections) + Config.c2 * 0 + Config.c3 * 0
            return 0
        else :
            if(not genome.connections) :
                return Config.c1 * 0 + Config.c2 * len(self.connections) + Config.c3 * 0

        all_connections = [connection for _,connection in self.connections.items()]
        all_connections += [connection for _,connection in genome.connections.items()]

        N = max(len(self.connections),len(genome.connections))
        min_innovation = min(max(self.connections),max(genome.connections))

        matching = 0
        average_weights = 0
        excess = 0
        disjoint = 0

        for innovation,connection in self.connections.items() :
            if(innovation in genome.connections) :
                matching += 1
                average_weights += abs(connection.weight - genome.connections[innovation].weight)
            elif(innovation < min_innovation) :
                disjoint += 1
            else :
                excess += 1

        for innovation,connection in genome.connections.items() :
            if(innovation in self.connections) :
                matching += 1
                average_weights += abs(connection.weight - self.connections[innovation].weight)
            elif(innovation < min_innovation) :
                disjoint += 1
            else :
                excess += 1

        if(matching != 0) :
            average_weights /= matching

        return (Config.c1 * disjoint) / N + (Config.c2 * excess) / N + Config.c3 * average_weights

    @staticmethod
    def generateRandomBaseGenome(input_size,output_size) :
        input_nodes = {}
        for i in range(input_size) :
            id = Config.getNodeId()
            activation = random.choice(list(Activation.methods.items()))[1]
            node = Node('input',activation,id)
            node.order = i
            input_nodes[id] = node

        output_nodes = {}
        for i in range(output_size) :
            id = Config.getNodeId()
            activation = random.choice(list(Activation.methods.items()))[1]
            node = Node('output',activation,id)
            node.order = i
            output_nodes[id] = node

        connections = {}
        for input_id,input_node in input_nodes.items() :
            for output_id,output_node in output_nodes.items() :
                if(np.random.uniform() <= 0.5) :
                    weight = random.uniform(-1,1)
                    innovation = Config.getInnovationId()
                    if(np.random.uniform() <= 0.5) :
                        enabled = True
                    else :
                        enabled = False
                    connection = Connection(input_node,output_node,innovation,weight,enabled)
                    connections[innovation] = connection

        nodes = {}
        for id,node in input_nodes.items() :
            nodes[id] = node
        for id,node in output_nodes.items() :
            nodes[id] = node
        genome = Genome(nodes,connections,Config.getGenomeId())
        return genome

    @staticmethod
    def generateNewGenome(parent_1,parent_2) :
        if(parent_1.fitness > parent_2.fitness) :
            fittest_parent = parent_1
            other_parent = parent_2
        elif(parent_1.fitness < parent_2.fitness) :
            fittest_parent = parent_2
            other_parent = parent_1
        elif(np.random.uniform() <= 0.5) :
            fittest_parent = parent_1
            other_parent = parent_2
        else :
            fittest_parent = parent_2
            other_parent = parent_1

        fittest_connections = fittest_parent.connections
        other_connections = other_parent.connections

        child_nodes = {}
        child_connections = {}

        for key,fittest_connection in fittest_connections.items() :
            chosen_connection = None
            if(key in other_connections) :
                Genome.__handle_matching_connetions(fittest_connection,other_connections[key],child_nodes,child_connections)

            else :
                Genome.__handle_disjoint_or_excess(fittest_connection,child_nodes,child_connections)

        for id,node in fittest_parent.nodes.items() :
            if(id not in child_nodes) :
                child_nodes[id] = node.clone()

        Genome.__mutate(child_nodes,child_connections)

        child = Genome(child_nodes,child_connections,Config.getGenomeId())

        return child

    @staticmethod
    def __handle_matching_connetions(fittest_connection,other_connection,child_nodes,child_connections) :
        if(np.random.uniform() <= 0.5) :
            child_connection = fittest_connection.clone()
            child_nodes[child_connection.input.id] = child_connection.input
            child_nodes[child_connection.output.id] = child_connection.output
        else :
            input_node = fittest_connection.input.clone()
            input_node.activation = other_connection.input.activation
            output_node = fittest_connection.output.clone()
            output_node.activation = other_connection.output.activation
            child_connection = Connection(input_node,output_node,other_connection.innovation,other_connection.weight,other_connection.enabled)
            child_nodes[input_node.id] = input_node
            child_nodes[output_node.id] = output_node

        child_connections[child_connection.innovation] = child_connection

    @staticmethod
    def __handle_disjoint_or_excess(fittest_connection,child_nodes,child_connections) :
        child_connection = fittest_connection.clone()
        child_nodes[child_connection.input.id] = child_connection.input
        child_nodes[child_connection.output.id] = child_connection.output
        child_connections[child_connection.innovation] = child_connection

    @staticmethod
    def __mutate(nodes,connections) :
        Genome.__mutation_weight(connections)
        Genome.__mutation_enabled(connections)
        Genome.__mutation_remove_connection(connections)
        Genome.__mutation_add_connection(nodes,connections)
        Genome.__mutation_remove_node(nodes,connections)
        Genome.__mutation_add_node(nodes,connections)

    @staticmethod
    def __mutation_activation(nodes) :
        if(np.random.uniform() <= Config.mutation_activation_rate) :
            node = random.choice(list(nodes.items()))[1]
            node.activation = random.choice(list(Activation.methods.items()))[1]

    @staticmethod
    def __mutation_weight(connections) :
        if(not connections) :
            return
        if(np.random.uniform() <= Config.mutation_weight_rate) :
            connection = random.choice(list(connections.items()))[1]
            connection.weight += random.uniform(Config.mutation_min,Config.mutation_max)

    @staticmethod
    def __mutation_enabled(connections) :
        if(not connections) :
            return
        if(np.random.uniform() <= Config.mutation_enabled_rate) :
            connection = random.choice(list(connections.items()))[1]
            if(connection.enabled) :
                connection.enabled = False
            else :
                connection.enabled = True

    @staticmethod
    def __mutation_add_connection(nodes,connections) :
        if(np.random.uniform() <= Config.mutation_new_connection_rate) :
            node_1 = nodes[random.choice(list(nodes.keys()))]

            candidates = {}
            for id,node in nodes.items() :
                if(node_1.type == node.type and node.type != 'hidden') :
                    continue
                if(node_1.id == node.id) :
                    continue
                candidates[id] = node

            for _,connection in connections.items() :
                input_node = connection.input
                output_node = connection.output
                if(node_1.id == input_node.id) :
                    node = output_node
                elif(node_1.id == output_node.id) :
                    node = input_node

                if(node.id in candidates) :
                    del candidates[node.id]

            if(candidates) :
                node_2 = nodes[random.choice(list(candidates.keys()))]
                innovation = Config.getInnovationId()
                if(node_1.type == 'input') :
                    new_connection = Connection(node_1,node_2,innovation,1,True)
                elif(node_1.type == 'output') :
                    new_connection = Connection(node_2,node_1,innovation,1,True)
                else :
                    new_connection = Connection(node_1,node_2,innovation,1,True)

                connections[innovation] = new_connection

                Genome.__fix_cycles(nodes,connections)

    @staticmethod
    def __fix_cycles(child_nodes,child_connections) :
        input_connections = Genome.__getNodesInputConnections(child_connections)
        for id,node in child_nodes.items() :
            connections_to_remove = Genome.__fix_cycles_helper(node,None,input_connections,{})
            for connection in connections_to_remove :
                if(connection.innovation in child_connections) :
                    del child_connections[connection.innovation]

    @staticmethod
    def __fix_cycles_helper(node,connection,input_connections,visited) :
        if(node.type == 'input' or node.id not in input_connections) :
            return []
        if(node.id in visited) :
            return [connection]
        visited[node.id] = node
        node_connections = input_connections[node.id]
        connections_to_remove = []
        for input_connection in node_connections :
            connections_to_remove += Genome.__fix_cycles_helper(input_connection.input,input_connection,input_connections,visited)

        return connections_to_remove

    @staticmethod
    def __mutation_remove_connection(connections) :
        if(not connections) :
            return
        if(np.random.uniform() <= Config.mutation_remove_connection_rate) :
            connection = random.choice(list(connections.items()))[1]
            del connections[connection.innovation]

    @staticmethod
    def __mutation_add_node(nodes,connections) :
        if(not connections) :
            return
        if(np.random.uniform() <= Config.mutation_new_node_rate) :
            connection = random.choice(list(connections.items()))[1]
            connection.enabled = False
            new_node = Node('hidden',connection.input.activation,Config.getNodeId())
            input_1 = connection.input
            new_connection_1 = Connection(input_1,new_node,innovation=Config.getInnovationId(),weight=1,enabled=True)
            output_2 = connection.output
            new_connection_2 = Connection(new_node,output_2,innovation=Config.getInnovationId(),weight=connection.weight,enabled=True)
            nodes[new_node.id] = new_node
            connections[new_connection_1.innovation] = new_connection_1
            connections[new_connection_2.innovation] = new_connection_2

    @staticmethod
    def __mutation_remove_node(nodes,connections) :
        if(np.random.uniform() <= Config.mutation_remove_node_rate) :
            hidden_nodes = [node.id for _,node in nodes.items() if node.type == 'hidden']
            if(len(hidden_nodes) != 0) :
                node_id = random.choice(hidden_nodes)
                del nodes[node_id]
                connections_to_remove = []
                for index,connection in connections.items() :
                    if(connection.input.id == node_id or connection.output.id == node_id) :
                        connections_to_remove.append(index)

                for connection in connections_to_remove :
                    del connections[connection]


import utime as time

class Control2:
    """! @brief The Control class carries out the the calculations for a PID controller based off of Kp, Ki, and Kd,
    and the set point (goal), and the current position
    """
    def __init__(self, kp, ki, kd, setpt):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpt = setpt

        self.I = 0
        self.error_prev = 0
        self.t_prev = time.ticks_ms()
        """! Initialization takes in the controller parameters and intiallizes the controller object
        @param kp: Kp is the proportional controller constant (must range from 0 --> 4 for this application)
        @param ki: Ki is the integral controller constant
        @param kd: Kd is the derivative controller constant
        @param setpt: setpt is the setpoint, or the encoder value for the "goal" position of the motor
        """

    def run(self, pos):
        """! The run method takes in the current motor position and calculates and scales the error to ouput the
            necessary motor effort for a PID controller.
        @param pos: pos is the current motor position as read by the encoder
        """
        # Calculate positional error
        self.error = self.setpt - pos
        self.t = time.ticks_ms()

        # Perform PID control and configure P,I, and D outputs
        self.psiP = self.kp*self.error
        self.psiI = self.I + self.ki*self.error*(self.t - self.t_prev)
        self.psiD = self.kd*(self.error - self.error_prev)/(self.t - self.t_prev)
        self.psi = self.psiP + self.psiI + self.psiD

        #print(f"error: {self.error:.2f}\teffort: {self.psi:.2f}\tP: {self.psiP:.2f}\tI: {self.psiI:.2f}\tD: {self.psiD:.2f}")

        self.error_prev = self.error
        self.t_prev = self.t

        return self.psi

    def set_setpoint(self, setpt_new):
        """! set_setpoint takes in the new setpoint, or goal position, of the motor and applies it to the controller object
        @param setpt_new: setpt_new is the new setpoint for the motor, given in an encoder position
        """
        self.setpt = setpt_new
        self.I = 0
        self.t = time.ticks_ms()
        return self.setpt

    def set_Kp(self, kp_new):
        """! set_Kp takes in the new Kp, proortional controller constant, and applies it to the controller object
        @param kp_new: kp_new is the new Kp
        """
        self.kp = kp_new
        print('kp is:', self.kp)
        return self.kp

if __name__ == '__main__':
    con = Control2(.15, .15, .15, 1000)


import math

with open('inputs/day06.txt', 'r') as file:
    memory = [int(n) for n in file.readline().split()]

SIZE = len(memory)

to_string = lambda a: '-'.join([str(n) for n in a])
seen_configs = set()
cycles = 0

def redistribute() -> str:
    max_element_index = max(enumerate(memory), key=lambda x: x[1])[0]
    starting_value = memory[max_element_index]
    distribution = int(math.ceil(starting_value / SIZE))

    memory[max_element_index] = 0
    i = max_element_index + 1

    while starting_value > 0:
        i %= SIZE
        memory[i] += min(distribution, starting_value)
        starting_value -= distribution
        i += 1

    return to_string(memory)

in_infinite_loop = False
loop_start = ''

while True:
    current_config = redistribute()

    if in_infinite_loop:
        cycles += 1
        if current_config == loop_start and cycles > 1:
            break
    else:
        if current_config in seen_configs:
            in_infinite_loop = True
            loop_start = current_config
        seen_configs.add(current_config)

result = cycles
print(f'Day 6 Part 2: The infinite loop is {result} cycles long.\n')


import requests
import selenium
from selenium import webdriver
from bs4 import BeautifulSoup

from pymongo import MongoClient
import datetime

client = MongoClient('localhost', 27017)
db = client.dbsparta
collection = db.collection
docs = collection.find({})
weekno = datetime.datetime.today().weekday()

def update():
    players = list(db.myfantasystar.find({}, {'_id': False}, {'position': False}, {'match': False}, {'goal': False},
                                         {'assi': False}, {'img_url': False}).sort('name', ))
    if weekno == 2:
        players['like'] = '0'
        players['pre_like'] = '0'
    else:
        players['pre_like'] = players['like']
        players['like'] = '0'

    print(players)

    # like_info = {
    #     'like': players['pre_like'],
    #     'pre_like': players['like'],
    # }
    # db.likedb.insert_one(like_info)

#
# def like_reset():
#     like = list(db.myfantasystar.find({}, {'_id': False}).sort('name', ))
#     pre_like =
#     if weekno == 2:
#         like = '0'
#         pre_like = '0'
#
# #오잉?! like를 0으로 선언했기 때문에 like를 바로 넣으면 안될듯?
#
# def like_update():
#     if weekno != 2:
#         pre_like = 'like'
#         like = '0'


import cv2 as cv
import mediapipe as mp
import numpy as np
import time
class poseDetector:
    def __init__(self, mode=False, upBody = False, smooth=True, detectionCon = 0.5, trackCon = 0.5):
        self.mode = mode
        self.upBody = upBody
        self.smooth=smooth
        self.detectionCon = detectionCon
        self.trackCon = trackCon
        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose()
        #self.pose = self.mpPose.Pose(self.mode, self.upBody, self.smooth, self.detectionCon, self.trackCon)
        self.mpDraw = mp.solutions.drawing_utils

    def findPose(self, img, draw=True):
        imgRGB = cv.cvtColor(img, cv.COLOR_BGR2RGB)
        self.results = self.pose.process(imgRGB)
        # print(self.results.pose_landmarks)
        if self.results.pose_landmarks:
            if draw:
                self.mpDraw.draw_landmarks(img, self.results.pose_landmarks, self.mpPose.POSE_CONNECTIONS)
        return img

    def findPosition(self, img, draw=False):
        lmlist = []
        if self.results.pose_landmarks:
            for id, lm in enumerate(self.results.pose_landmarks):
                h, w, c = img.shape
                cx, cy = int(lm.x * w), int(lm.y * h)
                lmlist.append([id, cx, cy])
                if draw:
                    cv.circle(img, (cx, cy), 8, (255, 0, 255), cv.FILLED)
        return lmlist

def main():
    cap = cv.VideoCapture('PoseVideos/KamariyaChoreography.mp4')
    detector = poseDetector()
    pTime = 0
    while True:
        success, img = cap.read()
        # img = cv.flip(img, 1)
        img = detector.findPose(img)
        # lmlist = detector.findPosition(img)
        # if len(lmlist) != 0:
        #     print(lmlist[4])
        cTime = time.time()
        fps = 1 / (cTime - pTime)
        pTime = cTime
        cv.putText(img, str(int(fps)), (10, 70), cv.FONT_HERSHEY_PLAIN, 3, (255, 0, 255), 3)
        cv.imshow("Video", img)
        key = cv.waitKey(1) & 0xFF
        if key == ord('q'):
            break
    cv.destroyAllWindows()

if __name__ == "__main__":
    main()


import glob
import os
import numpy as np
from planetaryimage import PDS3Image

# Start batch, end batch
startId = 3578
endId = 3584

for id in range(startId, endId):
    # create path if needed
    PDS_path = "CTX/" + str(id) + "/PDS3/"
    NPY_path = "CTX/" + str(id) + "/NPY/"
    if not os.path.exists(NPY_path):
        os.makedirs(NPY_path)

    # get all images
    files = glob.glob(PDS_path + "*.img")
    mapping = []
    for img in range(len(files)):
        # convert each image and save meta information
        print("Opening: ", img , " File:", files[img])
        image = PDS3Image.open(files[img])
        npy = image.image[:,:]
        print("Saving: ", img , " File:", files[img])
        np.save(NPY_path + str(img) + ".npy", npy)
        meta = {
            "id": img,
            "name": os.path.basename(files[img]),
        }
        mapping.append(meta)
    np.save("CTX/" + str(id) + "/mapping.npy", mapping)


directions = [[-1, 0], [0, 1], [1, 0], [0, -1]]

def numberOfIslands(matrix):
    if len(matrix) == 0:
        return 0
    counter = 0
    for row in range(len(matrix)):
        for col in range(len(matrix[0])):
            if matrix[row][col] == 1:
                counter += 1
                matrix[row][col] = 0
                queue = []
                queue.append([row, col])

                while len(queue):
                    currentPos = queue.pop(0)
                    currentRow = currentPos[0]
                    currentCol = currentPos[1]

                    for i in range(len(directions)):
                        currentDir = directions[i]
                        nextRow = currentRow + currentDir[0]
                        nextCol = currentCol + currentDir[1]

                        if nextRow < 0 or nextRow >= len(matrix) or nextCol < 0 or nextCol >= len(matrix[0]):
                            continue
                        if matrix[nextRow][nextCol] == 1:
                            queue.append([nextRow, nextCol])
                            matrix[nextRow][nextCol] = 0

    return counter


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from refinery.units import Unit

class cswap(Unit):
    """
    Swap the case of the input string; all lowercase letters are turned into their uppercase
    variant and vice-versa.
    """
    def process(self, data: bytearray):
        lcase = bytes(range(B'a'[0], B'z'[0] + 1))
        ucase = bytes(range(B'A'[0], B'Z'[0] + 1))
        delta = lcase[0] - ucase[0]
        for k, letter in enumerate(data):
            if letter in ucase:
                data[k] += delta
            elif letter in lcase:
                data[k] -= delta
        return data


import statistics

def part1(nums):
    c = statistics.median(nums)
    return sum(abs(x - c) for x in nums)

# simple gradient descent
def fmin(p0, f, num_it=1000):
    EPS = 0.001
    p = p0
    for i in range(num_it):
        pr = f(p + EPS)
        pl = f(p - EPS)
        dp = (pr - pl)/(2 * EPS)
        p = p - dp * EPS
    return p

def part2(nums):
    def f(x):
        res = 0
        for n in nums:
            d = abs(x - n)
            res += d * (d + 1)/2
        return res
    minx = int(round(fmin(statistics.median(nums), f)))
    return f(minx)

def solution():
    nums = [int(x) for x in open('../data/07.txt').read().split(',')]
    print(f"Answer 1: {int(part1(nums))}\nAnswer 2: {int(part2(nums))}")


vec = [11, 9, 7, 6, 5, 4, 3, 2, 1, 0, -1, - 3, -6, -6, -7, -7, -8, - 8]

#questão 1 e 2
def ordenar(vec:list, reverse = False):

    i = 1

    while(True):

        prev = vec[i - 1]
        cur = vec[i]

        #troca de maior que para menor que, se for reverse
        if(cur > prev if reverse else cur < prev):
            (vec[i - 1], vec[i]) = (cur, prev)
            i = max(1, i - 1)
        else:
            i += 1

        if(i > len(vec) - 1):
            break

    return vec

#questão 3
def mmax(vec:list):

    ma = vec[0]

    for cur in vec:
        if(cur > ma):
            ma = cur

    return ma

#questão 3
def mmin(vec:list):

    mi = vec[0]

    for cur in vec:
        if(cur < mi):
            mi = cur

    return mi

#questão 5
def remover_duplicados(vec:list):
    return list(set(vec))

#questão 4
def segundo_menor(vec:list):
    #ordena e depois pega o segundo menor pelo index 1 do array
    return ordenar(remover_duplicados(vec))[1]

#questão 7
def tercmax(vec:list):
    #ordena e depois pega o terceiro maior pelo index 2 do array
    return ordenar(remover_duplicados(vec), True)[2]

#questão 6
def par_impar(vec:list):

    resul = [0, 0]

    for cur in vec:
        resul[cur % 2] += 1

    print(f"o vetor tem {resul[0]} numeros pares, e {resul[1]} numeros impares")
    return resul

#questão 8
def mediana(vec:list):

    size = len(vec)
    index = int(size / 2) + 1

    orde = ordenar(vec)

    if(size % 2 == 0):
        return (orde[index] + orde[index - 1]) / 2
    else:
        return orde[index]

print("lista ordenada: ", ordenar(vec, False))
print("lista ordenada decrescente: ", ordenar(vec, True))
print("maior valor: " , mmax(vec))
print("menor valor: " , mmin(vec))
print("segundo menor: " , segundo_menor(vec))
print("sem duplicados: " , remover_duplicados(vec))
print("par e impar respectivamente: " , par_impar(vec))
print("terceiro maior: " , tercmax(vec))
print("mediana do vetor: " , mediana(vec))


"""
Temperature and Humidity and Barometric sensors
===============================================

"""

from rfxcom.protocol.base import BasePacketHandler
from rfxcom.protocol.rfxpacketutils import RfxPacketUtils

class TempHumidityBaro(BasePacketHandler):
    """
    ====    ====
    Byte    Meaning
    ====    ====
    0       Packet Length, 0x0D (excludes this byte)
    1       Packet Type, 0x54
    2       Sub Type
    3       Sequence Number
    4       ID 1
    5       ID 2
    6       Temperature High (7 bits), Temperature sign (1 bit)
    7       Temperature Low
    8       Humidity
    9       Humidity Status
    10      RSSI and Battery Level
    ====    ====

    """
    def __init__(self, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.PACKET_TYPES = {
            0x54: "Temperature and humidity and barometric sensors"
        }
        self.PACKET_SUBTYPES = {
            0x01: 'BTHR918',
            0x02: 'BTHR918N, BTHR968',
        }

    def parse(self, data):
        """Parse a 14 bytes packet in the TemperatureHumidity format and return a
        dictionary containing the data extracted. An example of a return value
        would be:

        .. code-block:: python

            {
                'id': "0x2EB2",
                'packet_length': 13,
                'packet_type': 84,
                'packet_type_name': 'Temperature and humidity and barometric \
                 sensors',
                'sequence_number': 0,
                'packet_subtype': 1,
                'packet_subtype_name': "BTHR918",
                'temperature': 21.3,
                'humidity': 91,
                'humidity_status': "Wet"
                'signal_level': 9,
                'battery_level': 6,
            }

        :param data: bytearray to be parsed
        :type data: bytearray

        :return: Data dictionary containing the parsed values
        :rtype: dict
        """

        self.validate_packet(data)

        id_ = self.dump_hex(data[4:6])
        temperature = ((data[6] & 0x7f) * 256 + data[7]) / 10
        signbit = data[6] & 0x80
        if signbit != 0:
            temperature = -temperature
        humidity = data[8]
        humidity_status = self._extract_humidity_status(data[9])
        baro = data[10] * 256 + data[11]
        forecast = data[12]
        forecast_status = self._extract_forecast(forecast)

        sensor_specific = {
            'id': id_,
            'temperature': temperature,
            'humidity': humidity,
            'humidity_status': humidity_status,
            'barometry': baro,
            'forecast': forecast,
            'forecast_status': forecast_status,
        }

        results = self.parse_header_part(data)
        results.update(RfxPacketUtils.parse_signal_and_battery(data[13]))
        results.update(sensor_specific)

        return results

    def _extract_humidity_status(self, data):
        """Extract the humidity status.

        :param data: byte to be parsed
        :type data: byte

        :return: String containing the human readable status
        :rtype: string
        """
        if data == 0x00:
            return "Dry"
        elif data == 0x01:
            return "Comfort"
        elif data == 0x02:
            return "Normal"
        elif data == 0x03:
            return "Wet"
        else:
            return "--??--"

    def _extract_forecast(self, data):
        """Extract the forecast info.

        :param data: byte to be parsed
        :type data: byte

        :return: String containing the human readable forecast
        :rtype: string
        """
        if data == 0x00:
            return "No forecast available"
        elif data == 0x01:
            return "Sunny"
        elif data == 0x02:
            return "Partly cloudy"
        elif data == 0x03:
            return "Cloudy"
        elif data == 0x04:
            return "Rainy"
        else:
            return "--??--"


"""
Dungeons & Dragons tool-kit
A module containing functions that simulate/calculate dice-rolls/scores

todo:
    - [x] code out rolls
        - [x] initiative
        - [x] attack
        - [x] damage
    - [x] integrate the probability of the score you rolled into:
        - [x] initiative
        - [x] attack
        - [x] damage
    - [] piercer feat
    - [] refactor damage roll as a class

"""
import statistics
from typing import *
import random
import functools
import itertools

import numpy
import pandas as pd

D20: int = 20

def roll_initiative() -> int:
    modifiers = {
        'init_bonus': 3,
        'widsom': 2,
    }
    roll = roll_vantage(D20, vantage='advantage')
    init = roll + sum(modifiers.values())
    print(f'init={init}, roll={roll}, {dict2str(modifiers)}')
    return init

def roll_attack(n: int = 2, vantage: Optional[str] = None) -> List:
    attacks = []
    for _ in range(n):
        modifiers = {
            'crossbow': 5,
            'fighting_style': 2,
        }
        if vantage:
            roll = roll_vantage(D20, vantage=vantage)
        else:
            roll = random.randint(1, D20)
        attack = roll + sum(modifiers.values())
        if roll == 20:
            print(f'CRITICAL HIT!!! - perfect {roll}')
        print(f'attack={attack}, roll={roll}, {dict2str(modifiers)}')
        attacks.append(attack)
    return attacks

def roll_damage(
        n: int = 1,
        d: int = 6,
        dread_ambush: bool = False,
        humanoid: bool = False,
        hunters_mark: bool = False,
        vantage: Optional[str] = None,
        critical_hit: Optional[Tuple[bool, ...]] = False,
) -> int:
    n = 3 if dread_ambush else n
    damages = []

    for i in range(n):
        modifers = {
            'crossbow': 3,
        }
        if humanoid: modifers['humanoid'] = 2
        if hunters_mark:
            modifers['hunters_mark'] = random.randint(1, 6)
            if critical_hit[i]:
                modifers['hunters_mark_critical'] = 6
        if dread_ambush and (i == 0):
            modifers['dread_ambush'] = random.randint(1, 8)
            if critical_hit[i]:
                modifers['dread_ambush_critical'] = 8

        if vantage:
            roll = roll_vantage(d, vantage=vantage)
        else:
            roll = random.randint(1, d)
            if critical_hit[i]:
                modifers['roll_critical'] = d
        damage = roll + sum(modifers.values())
        damages.append(damage)
        print(f'damage={damage}, roll={roll}, {dict2str(modifers)}')

    dread_ambush_roll = modifers.get('dread_ambush', 0)
    hunters_mark_roll = modifers.get('hunters_mark', 0)
    dread_ambush = dread_ambush_roll - 4
    hunters_mark = hunters_mark_roll - 3
    roll = roll - 3
    # whichever is lowest roll again

    total_damage = sum(damages)
    print(f'total_damage: {total_damage}')
    return total_damage

def roll_vantage(d: int, vantage: str) -> int:
    rolls = sorted([random.randint(1, d) for _ in range(2)], reverse=True)
    func = {
        'advantage': max,
        'disadvantage': min,
    }[vantage]
    roll = func(rolls)
    print(f'{vantage}={rolls}')
    return roll

def prob_of_roll(dice: Tuple[int, ...], total: Optional[int] = None) -> pd.DataFrame:
    """return the score probability distribution of the passed dice. If total passed return the
    specfic probability corresponding to that total"""
    dice = {f'{i}_d{d}': range(1, d + 1) for i, d in enumerate(dice)}
    combinations = list(itertools.product(*dice.values()))
    df = (pd.DataFrame(combinations, columns=dice.keys())
          .assign(total=lambda x: x.sum(1))
          .assign(dist=lambda x: x.groupby('total').transform('count').iloc[:, 0])
          .assign(prob=lambda x: x['dist'] / x.shape[0])
          .drop(dice.keys(), 1)
          .drop_duplicates()
          .set_index('total')
          .sort_index()
          )
    if total:
        return df.loc[total]
    return df

def dict2str(modifiers):
    modifiers_str = ', '.join([f'{k}={v}' for k, v in modifiers.items()])
    return modifiers_str

def avg_roll(func: Callable) -> float:
    return statistics.mean([func() for _ in range(10_000)])

roll_initiative()

roll_attack(
    n=2,
    vantage=None
)

args = {
    'n': 1,
    'dread_ambush': True,
    'humanoid': True,
    'hunters_mark': True,
    'critical_hit': (True, False, False),
}
roll_damage(**args)

roll_damage_partial = functools.partial(roll_damage, **args)

prob_of_roll((6, 6))


import  random
import  math

# Írj eljárást, mely a paraméterében kap két egész számot
# ezen két egész szám közötti
# páros számok összegét számolja ki az eljárást

def elsoFeladat(szamEgy,szamKetto):
    index = szamEgy
    osszeg = 0
    while index <= szamKetto:
        if index % 2 == 0:
            print(index)
            osszeg += index
        index+=1
    return osszeg

    # parosSzam = 0
    # osszeguk = 0
    # while szamEgy < szamKetto:
    #    if szamEgy % 2 == 0:
    #        parosSzam += szamEgy
    #    szamEgy+=1
    #print(parosSzam)

# írj függvényt amely generál 10 db véletlen számot
# -10 és 10 között
# Számold meg hány negatív szűm van közötte
# a visszatérési érték negativ számok száma

def negativ():
    negativakSzama = 0
    index = 0
    while index < 10:
        randomSzam = math.floor(random.random() * (10+10)-10)
        if randomSzam < 0:
            negativakSzama += 1
        index += 1
    return  negativakSzama

def negativKetto():
    negativakSzama = 0
    index = 0
    for i in range(0,20,1):
        randomSzam = math.floor(random.random() * (10 + 10) - 10)
        if randomSzam < 0:
            negativakSzama += 1
    return negativakSzama

def parosKetto(szamEgy,szamKetto):
    osszeg = 0
    for index in range(szamEgy,szamKetto,1):
        if index % 2 == 0:
            print(index)
            osszeg += index
    return osszeg

# Írj függvényt amely generál 10 db véletlen számot 12 és 33 között
# visszatér ezek összegével

def veletlenSzamokTizenkettoHarmincharom():
    osszeg = 0
    for i in range(0,10,1):
        randomSzam = math.floor(random.random()*(34-12)+12)
        osszeg += randomSzam
    return  osszeg

# Írj függvényt amely generál 8 db véletlen számot
# [20,50) intervalumban és visszatér ezek közül a legnagyobb számmal

def legnagyobbSzam():
    nagyobbSzam = 0
    for i in range(0,8,1):
        randomSzam = math.floor(random.random()*(50-20)+20)
        if nagyobbSzam < randomSzam:
            nagyobbSzam = randomSzam
    return  nagyobbSzam

# Kérjünk be 3 db egész számokat a felhasználótól
# mekkor a számok átlaga?

def bekeresHarom():
    osszegzes = 0
    atlag = 0
    for i in range(0,3,1):
        szamBekeres:int = int(input("Szeretnek kerni egy számot: "))
        osszegzes += szamBekeres
        atlag = round((osszegzes / 3))
    return atlag

# Kérjünk be egész számokat a felhasználótól amíg 0-at nem add
# Írjuk ki a számok átlagát

def bekeresAtlag():
    db = 0
    osszegzes = 0
    szamBekeres:int = int(input("Szeretnek kerni egész számokat, 0-ig: "))
    while not szamBekeres == 0:
        db += 1
        osszegzes += szamBekeres
        szamBekeres: int = int(input("Szeretnek kerni egész számokat, 0-ig: "))
    return round(osszegzes / db)


import NeuralNetwork as nn
import random as rand
import numpy as np
# import math
import loader
import PIL
import os

def getGuess(neuralNetwork, data):
    guesses = (np.asarray(neuralNetwork.guess(
        np.asarray(data[0] / 255, np.float64).flatten())).flatten())
    answers = np.asarray(data[1]).flatten()

    return [np.argmax(guesses),
            np.argmax(answers)]

def getGuessCustom(neuralNetwork, data):
    '''
    Given a neural network and a custom PIL image, return a value on what
    the network believes the image to be
    '''

    guesses = (np.asarray(neuralNetwork.guess(
        (np.asarray(list(data.getdata()), np.float64) / 255)
        .flatten())).flatten())

    print(guesses)

    return np.argmax(guesses)

def train(neuralNetwork, epochs, trainingData, decreaseLR=False):
    trainingDataLength = len(trainingData)

    for j in range(epochs):
        epochTraining = rand.sample(trainingData, trainingDataLength)

        for i in range(trainingDataLength):
            if(i % 10000 == 0):
                print('Iteration: ', i+(j*trainingDataLength))
            data = epochTraining[i]
            neuralNet.train(np.asarray(data[0] / 255, np.float64).flatten(),
                            np.asarray(data[1]).flatten())

        neuralNet.saveNetwork()
        print('Autosaved Network')

def predict(neuralNet, testingData):

    testingDataLength = len(testingData)

    correct = 0
    for i in range(testingDataLength):
        data = testingData[i]
        # rand.choice(testingData)
        guessesAnswersArray = getGuess(neuralNet, data)
        if(i % 100 == 0):
            print('Guess ', i+1, ': ', guessesAnswersArray[0],
                  ' Real: ', guessesAnswersArray[1])
        if(guessesAnswersArray[0] == guessesAnswersArray[1]):
            correct += 1

    print('Percentage of Correct Guesses: ',
          (correct/testingDataLength)*100, '% (', correct, ' out of ',
          testingDataLength, ')')
    print('Current learning rate: ', neuralNet.getLearningRate())

def predictCustom(neuralNet, testingData):
    guess = getGuessCustom(neuralNet, testingData)
    print('Guess :', guess)

    return guess

if __name__ == '__main__':

    dir = os.path.dirname(__file__)
    customFile = os.path.join(dir, 'custom')
    testIMG = PIL.Image.open(os.path.join(customFile, 'ci1.png'))

    data = loader.load_data()
    trainingData = data[0]
    testingData = data[1]

    neuralNet = nn.NeuralNetwork(784, 100, 100, 10, 0.001, True)

    epochs = 10

    # train(neuralNet, epochs, trainingData)

    predict(neuralNet, testingData)

    print('Saving Neural Network')
    neuralNet.saveNetwork()

B351 Intro to AI/Final Project/351FinalProject/HDRFullyConnected.py


import math as m
import os
import random

import cv2
import numpy as np

class Watermark:
    def __init__(self,
                 wt_path) -> None:
        self.pokes = self.loadWatermark(wt_path)
        pass

    def location(self, w, h, dx, dy):
        y, x = random.randint(0, w-dy-1), random.randint(0, h-dx-1)
        return y, x

    def embedding(self, input):
        wt = random.choice(self.pokes)

        w, h, _ = input.shape

        img_area = w*h
        mask_area = random.randint(10, 20) * img_area / 100

        new_size = int(m.sqrt(mask_area))

        wt = cv2.resize(wt, (new_size, new_size))

        dx, dy, _ = wt.shape

        mask = cv2.cvtColor(wt, cv2.COLOR_BGR2GRAY)

        mask = np.where(mask > 240, 0, 255)

        cv2.imwrite("mask.png", mask)

        y, x = self.location(w, h, dx, dy)

        # print(w, h, x, y, dx, dy)

        for i in range(3):
            print(input.shape, y+dx, x+dy)
            tmp = input[y:y+dx, x:x+dy, i]
            new_img_np = np.where(mask == 0, tmp, wt[:, :, i])
            input[y:y+dx, x:x+dy, i] = new_img_np
        output = input

        out_mask = np.zeros(input.shape, dtype=np.uint8)
        out_mask = cv2.cvtColor(out_mask, cv2.COLOR_BGR2GRAY)
        out_mask[y:y+dx, x:x+dy] = mask

        print(out_mask.shape)
        return output, out_mask

    def loadImage(self):
        pass

    def loadWatermark(self, wt_path):
        wtms = os.listdir(wt_path)

        marks = []

        for item in wtms:
            marks.append(cv2.imread(os.path.join(wt_path, item)))

        print(len(marks))

        return marks

def gen_data(data_path, save_path, label_path, watermark):

    images = os.listdir(data_path)

    for img_name in images:
        print(img_name)
        img = cv2.imread(os.path.join(data_path, img_name))

        new_image, label = watermark.embedding(img)

        cv2.imwrite(os.path.join(save_path, img_name), new_image)
        cv2.imwrite(os.path.join(label_path, img_name), label)

    pass

if __name__ == "__main__":

    wt_path = "/home/xuan/BKU/Thesis/data/watermark"
    data_path = "/home/xuan/BKU/Thesis/data/sample"
    save_path = "/home/xuan/BKU/Thesis/data/result"
    label_path = "/home/xuan/BKU/Thesis/data/label"
    watermark = Watermark(wt_path)

    gen_data(data_path, save_path, label_path, watermark)


from allauth.account.forms import SignupForm
from django.contrib.auth.models import Group
from django.core.mail import send_mail, EmailMultiAlternatives, mail_managers, mail_admins
from django import forms
# from .models import ContSub

# class ContactForm(forms.ModelForm):
#     class Meta:
#         model = ContSub
#         fields = '__all__'

# class CustomSignupForm(SignupForm):
#     def save(self, request):
#         user = super().save(request)
#         common_users = Group.objects.get(name="common users")
#         user.groups.add(common_users)
#         send_mail(
#             subject='Добро пожаловать в наш новостной сайт!',
#             message=f'{user.username}, вы успешно зарегистрировались!',
#             from_email=None,
#             recipient_list=[user.email],
#         )
#         return user

class CustomSignupForm(SignupForm):
    # Добавление пользователя в группу common_users
    def save(self, request):
        user = super().save(request)
        common_users = Group.objects.get(name="common users")
        user.groups.add(common_users)

        # Заголовок и текст письма
        subject = 'Добро пожаловать на новостной сайт.'
        text = f'{user.username}, вы успешно зарегистрировались на сайте!'
        html = (
            f'{user.username}, вы успешно зарегистрировались на '
            f'сайте My project!'
        )
        msg = EmailMultiAlternatives(
            subject=subject, body=text, from_email=None, to=[user.email]
        )
        msg.attach_alternative(html, "text/html")
        msg.send()

        # Рассылка писем менеджерам сайта и админам
        mail_managers(
            subject='Новый пользователь!',
            message=f'Пользователь {user.username} зарегистрировался на сайте.'
        )
        mail_admins(
            subject='Новый пользователь!',
            message=f'Пользователь {user.username} зарегистрировался на сайте.'
        )
        return user


# Import requests and Python's core JSON library.
import requests
import json

BASE_URL = 'http://pokeapi.co'

def query_pokeapi(resource_url):
    url = '{0}{1}'.format(BASE_URL, resource_url)
    response = requests.get(url)

    if response.status_code == 200:
        return json.loads(response.text)
    return None

def pokedex_description(pokemon_species):
    url = '{0}{1}'.format(BASE_URL, pokemon_species)
    response_two = requests.get(url)

    if response_two.status_code == 200:
        return json.loads(response_two.text)
    return None

charizard = query_pokeapi('/api/v1/pokemon/charizard/')
charizard_description = pokedex_description('/api/v2/pokemon-species/6/')

description = charizard_description['flavor_text_entries'][1]['flavor_text']
name = charizard['name']
sprite = charizard['sprites']['front_default']

# sprite = query_pokeapi(sprite_uri)
# print (sprite)
print (charizard['game_indices'][0]['game_index'])

# Create url containing url to the PokeAPI
# Make HTTP Request using requests' get() function
# Store results in a variable called response
# url = 'http://pokeapi.co/api/v1/pokemon/charizard/'
# url_two = 'https://pokeapi.co/api/v2/pokemon-species/6/'
# response = requests.get(url)

# response_two = requests.get(url_two)
# if response.status_code == 200:
#     data = json.loads(response.text)
#     data_two = json.loads(response_two.text)
#     print (data['name'])
#     print (data_two['flavor_text_entries'][1]['flavor_text'])
#     print (data['sprites']['front_default'])
#     print (data['game_indices'][0]['game_index'])
# else:
#     print ('An error occurred querying the API')