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)


'''
Create game (Online, Local with bots)
Connect to game
Settings
Exit
'''

import socket
import curses

from ncursesui.Elements import Button, List, Menu, Separator, TextField, UIElement, Window
from ncursesui.Utility import draw_borders, message_box
import draft_data

from mtgsdk import Card
from Networking import receive_msg, send_msg

DEFAULT_ADDRESS = 'localhost'
DEFAULT_PORT = str(draft_data.PORT)
ADDRESS_MAX_LENGTH = 20
PORT_MAX_LENGTH = 7
CARDS_LIST_HEIGHT = 20
CARDS_LIST_WIDTH = 20

class DraftClient:
    def __init__(self) -> None:
        pass

class DraftingMenu(Menu):
    def __init__(self, parent: Window, socket: socket.socket):
        draft_info = receive_msg(socket).split(' ')
        draft_name = draft_info[0]
        super().__init__(parent, 'Draft: {}'.format(draft_name.replace('_', ' ')))
        self.socket = socket
        self.pack_n = 0
        self.cards = []
        self.pile = []
        self.initUI()
        self.notify_about_next_pack()
        self.receive_msg()

    def initUI(self):
        def cards_list_element_click(choice: int, cursor: int, option: str):
            if message_box(self.parent, f'Take {option}?', ['No', 'Yes']) == 'Yes':
                # add card to pile and send it's mid to server
                self.pile += [self.cards[choice]]
                send_msg(self.socket, self.cards[choice].multiverseid)
                self.receive_msg()
        self.cards_list = List(self, [], CARDS_LIST_HEIGHT, CARDS_LIST_WIDTH, cards_list_element_click)
        self.cards_list.scroll_down_key = curses.KEY_DOWN
        self.cards_list.scroll_up_key = curses.KEY_UP
        self.cards_list.set_pos(0, 0)
        self.cards_list.set_focused(True)
        self.add_element(self.cards_list)

    def receive_msg(self):
        msg = receive_msg(self.socket)
        message_box(self.parent, msg)
        if msg == draft_data.NEXT_PACK:
            self.notify_about_next_pack()
            self.receive_msg()
            return
        if msg == draft_data.STOP_DRAFTING:
            self.stop_drafting()
            return
        # message is a list of mids
        mids = msg.split(' ')
        self.cards = [Card.from_id(mid) for mid in mids]
        self.cards_list.set_options([card.name for card in self.cards])

    def stop_drafting(self):
        # TODO
        message_box(self.parent, 'ENDED DRAFTING')
        self.socket.close()
        pass

    def notify_about_next_pack(self):
        self.pack_n += 1
        # self.parent.draw_window_with_message(f'PACK #{self.pack_n}')

class ClientWindow(Window):
    def __init__(self, window):
        super().__init__(window)

    def initUI(self):
        main_menu = Menu(self, 'Main menu')

        altext = 'Address: '
        address_label = UIElement(self, altext)
        address_label.set_pos(0, 0)

        address_text_field = TextField(self, DEFAULT_ADDRESS, ADDRESS_MAX_LENGTH)
        address_text_field.set_pos(0, len(altext))

        pltext = 'Port: '
        port_label = UIElement(self, pltext)
        port_label.set_pos(1, 0)

        port_text_field = TextField(self, DEFAULT_PORT, PORT_MAX_LENGTH)
        port_text_field.set_pos(1, len(altext))

        def connect_button_click():
            self.draw_window_with_message('Connecting to game...')
            address = address_text_field.text
            if address == 'localhost':
                address = socket.gethostname()
            try:
                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                s.connect((address, int(port_text_field.text)))
                self.draw_window_with_message('Queued for the game, waiting for other players...')
                self.current_menu = DraftingMenu(self, s)
            except (ConnectionRefusedError, socket.gaierror):
                self.draw_window_with_message('Could not connect to game!')
                self.window.getch()
            except ValueError:
                self.draw_window_with_message('Port has to be a number!')
                self.window.getch()

        connect_button = Button(self, 'Connect', connect_button_click)
        connect_button.set_pos(3, 0)
        connect_button.set_focused(True)

        address_text_field.prev = connect_button
        address_text_field.next = port_text_field

        port_text_field.prev = address_text_field
        port_text_field.next = connect_button

        connect_button.prev = port_text_field
        connect_button.next = address_text_field

        main_menu.add_element(address_label)
        main_menu.add_element(address_text_field)
        main_menu.add_element(port_label)
        main_menu.add_element(port_text_field)
        main_menu.add_element(Separator(self, 2))
        main_menu.add_element(connect_button)

        self.current_menu = main_menu

    def draw_window_with_message(self, message: str):
        height = 3
        width = len(message) + 2
        y = (self.HEIGHT - height) // 2
        x = (self.WIDTH - width) // 2
        window = curses.newwin(height, width, y, x)
        draw_borders(window, 'red-black')
        window.addstr(1, 1, message)
        window.refresh()

def main(stdscr):
    curses.curs_set(0)
    window = ClientWindow(stdscr)
    window.start()

curses.wrapper(main)


# McSimps Titanfall Map Exporter Tool
# Website: https://will.io/
# Modded(butchered) to output ply by Warmist
#   disclaimer: i don't know python...

import struct
from enum import Enum
import os
import math

#settings
map_name = 'mp_wargames'
map_unpacked_path='G:\\tmp\\mp_wargames\\' #path to unpacked bsp folder with "models" and "maps" subfolder
common_model_path_prefix='G:\\tmp\\mp_common\\' # path to unpacked englishclient_mp_common.bsp.pak000_dir.vpk (only models folder is needed)
#not settings
model_path_prefix=map_unpacked_path
map_path_prefix=map_unpacked_path+'\\maps\\' # path to unpacked bsp
dump_base = map_path_prefix+"out\\"
map_path =  map_path_prefix + map_name + '.bsp'

#mprt has bad angles, it might be wrongly rotated, but left it for future generations to admire...
mprt_path=""

def read_null_string(f):
    chars = []
    while True:
        c = f.read(1).decode('ascii')
        if c == chr(0) or c=='':
            return ''.join(chars)
        chars.append(c)

class LumpElement:
    @staticmethod
    def get_size():
        raise NotImplementedError()

class TextureData(LumpElement):
    def __init__(self, data):
        self.string_table_index = struct.unpack_from('<I', data, 12)[0]

    @staticmethod
    def get_size():
        return 36

class BumpLitVertex(LumpElement):
    def __init__(self, data):
        self.vertex_pos_index = struct.unpack_from('<I', data, 0)[0]
        self.vertex_normal_index = struct.unpack_from('<I', data, 4)[0]
        self.texcoord0 = struct.unpack_from('<ff', data, 8) # coord into albedo, normal, gloss, spec
        self.texcoord5 = struct.unpack_from('<ff', data, 20) # coord into lightmap

    @staticmethod
    def get_size():
        return 44

class UnlitVertex(LumpElement):
    def __init__(self, data):
        self.vertex_pos_index = struct.unpack_from('<I', data, 0)[0]
        self.vertex_normal_index = struct.unpack_from('<I', data, 4)[0]
        self.texcoord0 = struct.unpack_from('<ff', data, 8)

    @staticmethod
    def get_size():
        return 20

class UnlitTSVertex(LumpElement):
    def __init__(self, data):
        self.vertex_pos_index = struct.unpack_from('<I', data, 0)[0]
        self.vertex_normal_index = struct.unpack_from('<I', data, 4)[0]
        self.texcoord0 = struct.unpack_from('<ff', data, 8)

    @staticmethod
    def get_size():
        return 28

class MaterialSortElement(LumpElement):
    def __init__(self, data):
        self.texture_index = struct.unpack_from('<H', data, 0)[0]
        self.vertex_start_index = struct.unpack_from('<I', data, 8)[0]

    @staticmethod
    def get_size():
        return 12

class VertexType(Enum):
    LIT_FLAT = 0
    UNLIT = 1
    LIT_BUMP = 2
    UNLIT_TS = 3

class MeshElement(LumpElement):
    def __init__(self, data):
        self.indices_start_index = struct.unpack_from('<I', data, 0)[0]
        self.num_triangles = struct.unpack_from('<H', data, 4)[0]
        self.material_sort_index = struct.unpack_from('<H', data, 22)[0]
        self.flags = struct.unpack_from('<I', data, 24)[0]

    def get_vertex_type(self):
        temp = 0
        if self.flags & 0x400:
            temp |= 1
        if self.flags & 0x200:
            temp |= 2
        return VertexType(temp)

    @staticmethod
    def get_size():
        return 28

class PhyHeader(LumpElement): # Header in SourceIO
    def __init__(self, data):
        self.size = struct.unpack_from('<I', data, 0)[0]
        self.id = struct.unpack_from('<I', data, 4)[0]
        self.solidCount = struct.unpack_from('<I', data, 8)[0]
        self.checkSum = struct.unpack_from('<I', data, 12)[0]

    @staticmethod
    def get_size():
        return 16

class CompactSurfHeader(LumpElement):
    def __init__(self, data):
        self.size = struct.unpack_from('<I', data, 0)[0]
        self.id = struct.unpack_from('<I', data, 4)[0]
        self.version = struct.unpack_from('<H', data, 8)[0]
        self.modelType = struct.unpack_from('<H', data, 10)[0]
        self.surfaceSize=struct.unpack_from('<I', data, 12)[0]
        self.dragAxisAreas=struct.unpack_from('<3f', data, 16)
        self.axisMapSize=struct.unpack_from('<I', data, 28)[0]
        #collision model here
        self.cm_values=struct.unpack_from('<7f', data, 32)
        self.surface,self.offset_tree,*self.pad=struct.unpack_from('<4I',data,60)
        #some unk stuff here
        self.IPVS=struct.unpack_from('<I', data, 76)[0]

    @staticmethod
    def get_size():
        return 80
    @staticmethod
    def get_model_start():
        return 32

class PhyFaceHeader(LumpElement): #SourceIO ConvexLeaf
    def __init__(self, data):
        self.vert_offset = struct.unpack_from('<I', data, 0)[0] #vertex offset
        self.boneidx = struct.unpack_from('<I', data, 4)[0]
        self.unk=struct.unpack_from('<I', data, 8)[0]
        self.tri_count=struct.unpack_from('<h', data, 12)[0]
        self.pad=struct.unpack_from('<h', data, 14)[0]

    @staticmethod
    def get_size():
        return 16

class ConvexLeaf():
    def __init__(self):
        self.head: PhyHeader=None
        self.tris = []
        self.vertices=[]

class TreeNode():
    def __init__(self, data):
        self.right_node_offset, self.convex_offset, *self.floats = struct.unpack_from('2i5f',data,0)

        self.left_node: Optional[TreeNode] = None
        self.right_node: Optional[TreeNode] = None
        self.convex_leaf: Optional[ConvexLeaf] = None

    @property
    def is_leaf(self):
        return self.right_node_offset == 0

    @staticmethod
    def get_size():
        return 28

class location_saved(object):
    def __init__(self, f):
        self.f=f

    def __enter__(self):
        self.location = self.f.tell()
        return self.f

    def __exit__(self,exc_type, exc_value, exc_traceback):
        self.f.seek(self.location)

def load_leaf(f,all_indices):
    cv=ConvexLeaf()
    entry=f.tell()
    cv.entry=entry
    cv.head=PhyFaceHeader(f.read(PhyFaceHeader.get_size()))
    with location_saved(f):
        #all_indices = set()
        triangles = []

        for _ in range(cv.head.tri_count):
            f.seek(4,1)# pad
            tri=[]
            for i in range(3):
                tri.append(struct.unpack_from('<HH', f.read(4),0)[0])
            triangles.append(tri)
            all_indices.update(tri)
        #print(min(all_indices),max(all_indices))
        cv.triangles=triangles

        #min_tri_id=min(all_indices)
        #mapping={}
        #for i,id in zip(sorted(all_indices),range(max(all_indices)-min(all_indices)+1)):
        #    print(id,i)
        #    mapping[i]=id
        #cv.triangles[:]=map(lambda x: [x[0]-min_tri_id,x[1]-min_tri_id,x[2]-min_tri_id],triangles)
        #cv.triangles[:]=map(lambda x: [mapping[x[0]],mapping[x[1]],mapping[x[2]]],triangles)
        #f.seek(entry + cv.head.vert_offset)
        #vf=[]
        #for i in range(max(all_indices)-min_tri_id+1):
        #    vf=struct.unpack_from('<ffff',f.read(4*4))
        #    if (i+min_tri_id) in all_indices:
        #        cv.vertices.append([vf[0],vf[2],vf[1]])
        #print("Tris: {} l:{}".format(len(cv.triangles),len(all_indices)))
        return cv

def load_root_verts(f,n,count):
    verts=[]
    with location_saved(f):
        f.seek(n.convex_leaf.entry+n.convex_leaf.head.vert_offset)
        for i in range(count):
            vf=struct.unpack_from('<ffff',f.read(4*4))
            verts.append([vf[0],vf[2],vf[1]])
    return verts
def load_node(f,all_indices):
    entry=f.tell()
    #print("Node: {}".format(entry))
    n=TreeNode(f.read(TreeNode.get_size()))
    n.entry=entry
    #print("R: {} C: {} ".format(n.right_node_offset,n.convex_offset))
    with location_saved(f):
        #if n.is_leaf:
        if n.convex_offset!=0:
            with location_saved(f):
                f.seek(entry+n.convex_offset)
                n.convex_leaf=load_leaf(f,all_indices)
            #return n
            if n.is_leaf:
                return n
        n.left_node=load_node(f,all_indices)
        with location_saved(f):
            f.seek(entry+n.right_node_offset)
            n.right_node=load_node(f,all_indices)
    return n
def load_tree(f,root_offset):
    print("Loading with offset:{}".format(root_offset))
    with location_saved(f):
        f.seek(root_offset,1)
        all_indices=set()
        root=load_node(f,all_indices)
        verts=load_root_verts(f,root,len(all_indices))
        return root,verts

def concat_tree(n,tbl_tris):
    if n.is_leaf:
        #tbl_tris.extend(map(lambda x: [x[0]+cur_verts,x[1]+cur_verts,x[2]+cur_verts],n.convex_leaf.triangles))
        tbl_tris.extend(n.convex_leaf.triangles)
    else:
        if not(n.left_node is None):
            concat_tree(n.left_node,tbl_tris)
        if not(n.right_node is None):
            concat_tree(n.right_node,tbl_tris)

class MDL_Mesh():
    def __init__(self):
        self.tris=[]
        self.verts=[]
    def get_verts_transformed(self,origin,angles,scale):
        ca=math.cos(angles[0])
        cb=math.cos(angles[1])
        cc=math.cos(angles[2])

        sa=math.sin(angles[0])
        sb=math.sin(angles[1])
        sc=math.sin(angles[2])

        Axx = ca*cb
        Axy = ca*sb*sc - sa*cc
        Axz = ca*sb*cc + sa*sc

        Ayx = sa*cb;
        Ayy = sa*sb*sc + ca*cc;
        Ayz = sa*sb*cc - ca*sc;

        Azx = -sb;
        Azy = cb*sc;
        Azz = cb*cc;
        verts=[]
        for v in self.verts:
            nv=[0,0,0]
            nv[0]=(Axx*v[0]+Axy*v[1]+Axz*v[2])*scale+origin[0]
            nv[1]=(Ayx*v[0]+Ayy*v[1]+Ayz*v[2])*scale+origin[1]
            nv[2]=-(Azx*v[0]+Azy*v[1]+Azz*v[2])*scale+origin[2]
            verts.append(nv)
        return verts
    def get_tris_offset(self,offset):
        ret=[]
        for t in self.tris:
            ret.append([t[0]+offset,t[2]+offset,t[1]+offset])
        return ret

def dump_ply(verts,tris,fname):
    with open(fname, 'wb') as f:
            f.write(b"ply\nformat binary_little_endian 1.0\n");
            v_str="element vertex {}\n".format(len(verts))
            f.write(bytes(v_str,encoding='utf8'))
            f.write(b"property float x\nproperty float y\nproperty float z\n")
            f_str="element face {}\n".format(len(tris))
            f.write(bytes(f_str,encoding='utf8'))
            f.write(b"property list uchar int vertex_index\nend_header\n")
            for v in verts:
                f.write(struct.pack("<fff",v[0],v[1],v[2]))

            for t in tris:
                f.write(struct.pack("<cIII",b'\x03',t[0], t[1], t[2]))

mdl_cache={}
def load_mdl_file(path):
    ret_mesh=MDL_Mesh()
    with open(path,"rb") as f:
        print("Loading:",path)
        f.seek(0x6c)
        bbox_min=struct.unpack_from('<fff', f.read(12),0)
        bbox_max=struct.unpack_from('<fff', f.read(12),0)
        #print(bbox_max,bbox_min)
        #print((bbox_max[0]-bbox_min[0],bbox_max[1]-bbox_min[1],bbox_max[2]-bbox_min[2]))
        f.seek(0x1b8) # seek to phyOffset
        phy_offset=struct.unpack_from('<I', f.read(4), 0)[0]
        if phy_offset==0: #no phy, don't add the file
            return
        #print(path,phy_offset)
        f.seek(phy_offset)
        head0=PhyHeader(f.read(PhyHeader.get_size()))
        #print("Solid count:",head0.solidCount)
        next_header=phy_offset+PhyHeader.get_size()
        vertex_pos=next_header+head0.size
        max_vert=0
        bbox_phy_min=[99999,99999,99999]
        bbox_phy_max=[-99999,-99999,-99999]
        for i in range(head0.solidCount):
            assert(i==0)
            f.seek(next_header)
            chead=CompactSurfHeader(f.read(CompactSurfHeader.get_size()))
            next_header+=chead.size
            root_node,root_verts=load_tree(f,chead.offset_tree-CompactSurfHeader.get_size()+CompactSurfHeader.get_model_start())
            tmp_tris=[]
            concat_tree(root_node,tmp_tris)
            ret_mesh.verts=root_verts
            ret_mesh.tris=tmp_tris
            dump_ply(root_verts,tmp_tris,"tmp.ply")
            cur_pos=phy_offset+PhyHeader.get_size()

        #print(len(ret_mesh.tris),len(ret_mesh.verts))
        #print(((bbox_phy_max[0]-bbox_phy_min[0])*40,(bbox_phy_max[1]-bbox_phy_min[1])*40,(bbox_phy_max[2]-bbox_phy_min[2])*40))
        #dump_ply(ret_mesh.verts,ret_mesh.tris,"out.ply")
        #quit()
        #print(ret_mesh)
        return ret_mesh

def load_or_get_mdl(path):
    if path in mdl_cache:
        return mdl_cache[path]
    fullpath=model_path_prefix+path
    mexist=os.path.exists(fullpath)
    if not mexist:
        fullpath=common_model_path_prefix+path
    mexist=os.path.exists(fullpath)
    if not mexist:
        print("File not found:",path)
        quit()
    #print("Loading mdl:",fullpath)
    #print("Exists:",mexist)
    mdl_cache[path]=load_mdl_file(fullpath)
    return mdl_cache[path]

mdl_verts=[]
mdl_tris=[]
load_props=True

def load_mprt_file():
    print("Loading mprt file...")
    with open(mprt_path+map_name+".mprt","rb") as f:
        f.seek(12)
        while True:
            #mdl_path=''.join(iter(lambda: f.read(1).decode('ascii'), '\x00'))
            mdl_path=read_null_string(f)
            if not mdl_path:
                break
            mdl=load_or_get_mdl(mdl_path)
            origin=struct.unpack_from('<fff',f.read(12),0)
            angles=struct.unpack_from('<fff',f.read(12),0)
            #print(angles)
            angles=(math.pi*((angles[0])/180),math.pi*((angles[2])/180),math.pi*((angles[1])/180))
            scale=struct.unpack_from('<f',f.read(4),0)[0]
            #print(origin,angles,scale)
            if not mdl: #no physics: skip
                continue
            #print("Loading:",mdl_path,scale)
            offset=len(mdl_verts)
            mdl_verts.extend(mdl.get_verts_transformed(origin,angles,scale*40))
            mdl_tris.extend(mdl.get_tris_offset(offset))

def load_prop_lump():
    print("Loading prop lump...")
    with open(map_path+".0023.bsp_lump",'rb') as f:
        f.seek(0x14)
        name_count=struct.unpack_from('<I',f.read(4),0)[0]
        model_names = struct.iter_unpack("128s", f.read(128 * name_count))
        model_names = [t[0].replace(b"\0", b"").decode() for t in model_names]
        print("loaded {} prop names".format(name_count))
        prop_count=struct.unpack_from('<I',f.read(4),0)[0]
        f.seek(8,1)
        for i in range(prop_count):
            origin=struct.unpack_from('<fff',f.read(12),0)
            angles=struct.unpack_from('<fff',f.read(12),0)
            # 0 1 2 crane off axis!
            # 0 2 1 crane off axis!
            # 2 0 1 crane off axis!
            # 2 1 0 crane off axis!

            # 1 0 2

            # 1 2 0
            angles=(math.pi*((angles[1])/180),math.pi*((angles[2])/180),math.pi*((angles[0])/180))
            scale=struct.unpack_from('<f',f.read(4),0)[0]
            mdl_name_idx=struct.unpack_from('<H',f.read(2),0)[0]
            f.seek(34,1) # skip rest of struct
            #print(mdl_name_idx)
            #print("model_names[mdl_name_idx]:",model_names[mdl_name_idx])
            mdl=load_or_get_mdl(model_names[mdl_name_idx])
            if not mdl:
                continue
            offset=len(mdl_verts)
            mdl_verts.extend(mdl.get_verts_transformed(origin,angles,scale*40))
            mdl_tris.extend(mdl.get_tris_offset(offset))

if load_props:
    if mprt_path=="":
        load_prop_lump()
    else:
        load_mprt_file()
    #dump_ply(mdl_verts,mdl_tris,"out1.ply")
    print("Final mdl verts, tris:",len(mdl_verts),len(mdl_tris))

# Read all vertex position data
print("Reading vertex position data...")
with open(map_path + '.0003.bsp_lump', 'rb') as f:
    data = f.read()
    vertex_positions = [struct.unpack_from('<fff', data, i * 12) for i in range(len(data) // 12)]

# Read all vertex normals
print("Reading vertex normal data...")
with open(map_path + '.001e.bsp_lump', 'rb') as f:
    data = f.read()
    vertex_normals = [struct.unpack_from('<fff', data, i * 12) for i in range(len(data) // 12)]

# Read indices
print("Reading indices...")
with open(map_path + '.004f.bsp_lump', 'rb') as f:
    data = f.read()
    indices = [struct.unpack_from('<H', data, i * 2)[0] for i in range(len(data) // 2)]

# Read texture information
print("Reading texture information...")
with open(map_path + '.002c.bsp_lump', 'rb') as f:
    data = f.read()
    texture_string_offets = [struct.unpack_from('<I', data, i * 4)[0] for i in range(len(data) // 4)]

with open(map_path + '.002b.bsp_lump', 'rb') as f:
    texture_strings = []
    for offset in texture_string_offets:
        f.seek(offset)
        texture_strings.append(read_null_string(f))

textures = []
with open(map_path + '.0002.bsp_lump', 'rb') as f:
    data = f.read()
    elem_size = TextureData.get_size()
    for i in range(len(data) // elem_size):
        textures.append(TextureData(data[i*elem_size:(i+1)*elem_size]))

# Read bump lit vertices
print("Reading bump lit vertices...")
bump_lit_vertices = []
with open(map_path + '.0049.bsp_lump', 'rb') as f:
    data = f.read()
    elem_size = BumpLitVertex.get_size()
    for i in range(len(data) // elem_size):
        bump_lit_vertices.append(BumpLitVertex(data[i*elem_size:(i+1)*elem_size]))

# Read unlit vertices
print("Reading unlit vertices...")
unlit_vertices = []
unlit_vertex_bsp_lump = map_path + ".0047.bsp_lump"
if os.path.exists(unlit_vertex_bsp_lump):
    with open(unlit_vertex_bsp_lump, 'rb') as f:
        data = f.read()
        elem_size = UnlitVertex.get_size()
        for i in range(len(data) // elem_size):
            unlit_vertices.append(UnlitVertex(data[i*elem_size:(i+1)*elem_size]))

# Read unlit TS vertices
print("Reading unlit TS vertices...")
unlit_ts_vertices = []
unlit_ts_vertex_bsp_lump = map_path + ".004a.bsp_lump"
if os.path.exists(unlit_ts_vertex_bsp_lump ):
    with open(unlit_ts_vertex_bsp_lump , 'rb') as f:
      data = f.read()
      elem_size = UnlitTSVertex.get_size()
      for i in range(len(data) // elem_size):
          unlit_ts_vertices.append(UnlitTSVertex(data[i*elem_size:(i+1)*elem_size]))

vertex_arrays = [
    [],
    unlit_vertices,
    bump_lit_vertices,
    unlit_ts_vertices
]

# Read the material sort data
print("Reading material sort data...")
material_sorts = []
with open(map_path + '.0052.bsp_lump', 'rb') as f:
    data = f.read()
    elem_size = MaterialSortElement.get_size()
    for i in range(len(data) // elem_size):
        material_sorts.append(MaterialSortElement(data[i*elem_size:(i+1)*elem_size]))

# Read mesh information
print("Reading mesh data...")
meshes = []
with open(map_path + '.0050.bsp_lump', 'rb') as f:
    data = f.read()
    elem_size = MeshElement.get_size()
    for i in range(len(data) // elem_size):
        meshes.append(MeshElement(data[i*elem_size:(i+1)*elem_size]))

# Build combined model data
print("Building combined model data...")
combined_uvs = []
mesh_faces = []
texture_set = set()
for mesh_index in range(len(meshes)):
    faces = []
    mesh = meshes[mesh_index]
    if mesh.get_vertex_type().value==1 or mesh.get_vertex_type().value==3:
        continue
    mat = material_sorts[mesh.material_sort_index]
    texture_set.add(texture_strings[textures[mat.texture_index].string_table_index])
    for i in range(mesh.num_triangles * 3):
        vertex = vertex_arrays[mesh.get_vertex_type().value][mat.vertex_start_index + indices[mesh.indices_start_index + i]]
        combined_uvs.append(vertex.texcoord0)
        uv_idx = len(combined_uvs) - 1
        #faces.append((vertex.vertex_pos_index + 1, uv_idx + 1, vertex.vertex_normal_index + 1)) # obj files use 1 as index start
        faces.append((vertex.vertex_pos_index , uv_idx , vertex.vertex_normal_index )) #ply files use 0 as index start
    mesh_faces.append(faces)

# Build material files
# print('Building material files...')
# for i in range(len(textures)):
#     texture_string = texture_strings[textures[i].string_table_index]

#     # Work out the path to the actual texture
#     if os.path.isfile(dump_base + texture_string + '.png'):
#         path = dump_base + texture_string + '.png'
#     elif os.path.isfile(dump_base + texture_string + '_col.png'):
#         path = dump_base + texture_string + '_col.png'
#     else:
#         print('[!] Failed to find texture file for {}'.format(texture_string))
#         path = 'error.png'

    # # Write the material file
    # with open('{}\\tex{}.mtl'.format(map_name, i), 'w') as f:
    #      f.write('newmtl tex{}\n'.format(i))
    #      f.write('illum 1\n')
    #      f.write('Ka 1.0000 1.0000 1.0000\n')
    #      f.write('Kd 1.0000 1.0000 1.0000\n')
    #      f.write('map_Ka {}\n'.format(path))
    #      f.write('map_Kd {}\n'.format(path))

# Create obj file
print("Writing output file...")

# with open(map_name+".obj", 'w') as f:
#     f.write('o {}\n'.format(map_name))
#     for i in range(len(textures)):
#         f.write('mtllib tex{}.mtl\n'.format(i))

#     for v in vertex_positions:
#         f.write('v {} {} {}\n'.format(*v))

#     for v in vertex_normals:
#         f.write('vn {} {} {}\n'.format(*v))

#     for v in combined_uvs:
#         f.write('vt {} {}\n'.format(*v))

#     for i in range(len(mesh_faces)):
#         f.write('g {}\n'.format(i))
#         f.write('usemtl tex{}\n'.format(material_sorts[meshes[i].material_sort_index].texture_index))
#         faces = mesh_faces[i]
#         for i in range(len(faces) // 3):
#             f.write('f {}/{}/{} {}/{}/{} {}/{}/{}\n'.format(*faces[i*3], *faces[(i*3) + 1], *faces[(i*3) + 2]))

with open(map_name+".ply", 'wb') as f:
    f.write(b"ply\nformat binary_little_endian 1.0\n");
    v_str="element vertex {}\n".format(len(vertex_positions)+len(mdl_verts))
    #v_str="element vertex {}\n".format(len(mdl_verts))
    f.write(bytes(v_str,encoding='utf8'))
    f.write(b"property float x\nproperty float y\nproperty float z\n")

    num_triangles=0
    for mf in range(len(mesh_faces)):
        num_triangles+=len(mesh_faces[mf])//3

    f_str="element face {}\n".format(num_triangles+len(mdl_tris))
    #f_str="element face {}\n".format(len(mdl_tris))
    f.write(bytes(f_str,encoding='utf8'))
    f.write(b"property list uchar int vertex_index\nend_header\n")
    for v in vertex_positions:
        f.write(struct.pack("<fff",*v))

    for v in mdl_verts:
        f.write(struct.pack("<fff",*v))

    for i in range(len(mesh_faces)):
        faces = mesh_faces[i]
        for i in range(len(faces) // 3):
            f.write(struct.pack("<cIII",b'\x03',faces[i*3][0], faces[(i*3) + 1][0], faces[(i*3) + 2][0]))
    offset=len(vertex_positions)
    for t in mdl_tris:
        f.write(struct.pack("<cIII",b'\x03',t[0]+offset, t[1]+offset, t[2]+offset))


import json

from tqdm import tqdm

def build_and_save_data(
    abstract_file="ChemProt_Corpus/chemprot_training/chemprot_training_abstracts.tsv",
    entity_file="ChemProt_Corpus/chemprot_training/chemprot_training_entities.tsv",
    relation_file="ChemProt_Corpus/chemprot_training/chemprot_training_relations.tsv",
    target_file="chemprot_train.json",
):
    data = open(abstract_file).read().split("\n")
    id2text = {
        sent.split("\t", 1)[0]: sent.split("\t", 1)[1] for sent in data if len(sent) > 0
    }

    data = open(entity_file).read().split("\n")
    entities = {}
    for line in tqdm(data):
        items = line.split("\t", 5)
        if len(items) < 6:
            print(items)
            continue
        entities[f"{items[0]}.{items[1]}"] = items[-1]
    print(list(entities.items())[:10])

    ds = []

    data = open(relation_file).read().split("\n")
    for line in tqdm(data):
        items = line.split("\t", 5)
        if len(items) < 6:
            print(items)
            continue

        id = items[0]
        e1 = items[4].split(":", 1)[1]
        e2 = items[5].split(":", 1)[1]

        e1 = entities[f"{id}.{e1}"]
        e2 = entities[f"{id}.{e2}"]

        text = id2text[id]
        relation = items[3]

        sample = f"find relation between entities - {e1} and {e2} in {text}"
        target = relation

        ds.append({"sample": sample, "target": target})

    with open(target_file, "w") as f:
        json.dump(ds, f)

build_and_save_data(
    abstract_file="ChemProt_Corpus/chemprot_training/chemprot_training_abstracts.tsv",
    entity_file="ChemProt_Corpus/chemprot_training/chemprot_training_entities.tsv",
    relation_file="ChemProt_Corpus/chemprot_training/chemprot_training_relations.tsv",
    target_file="chemprot_train.json",
)
build_and_save_data(
    abstract_file="ChemProt_Corpus/chemprot_development/chemprot_development_abstracts.tsv",
    entity_file="ChemProt_Corpus/chemprot_development/chemprot_development_entities.tsv",
    relation_file="ChemProt_Corpus/chemprot_development/chemprot_development_relations.tsv",
    target_file="chemprot_valid.json",
)

from datasets import DatasetDict, load_dataset

ds = DatasetDict(
    {
        "train": load_dataset("json", data_files="chemprot_train.json", split="train"),
        "valid": load_dataset("json", data_files="chemprot_valid.json", split="train"),
    }
)

print(ds)

push_to_hub = False

if push_to_hub:
    ds.push_to_hub("ddp-iitm/chemprot", private=True)


"""
An Atom feed.
"""

from xml.etree import ElementTree
from datetime import datetime

# A puppy dies every time Fredrik Lundh serializes XML.
NS = "http://www.w3.org/2005/Atom"
ElementTree._namespace_map[NS] = "atom"

def qname(name):
    """
    Return a node name qualified with the Atom namespace.
    """
    return unicode(ElementTree.QName(NS, name))

class Atom(object):
    """
    An Atom feed.
    """

    # The maximum number of <entry> nodes.
    ENTRIES = 15

    def __init__(self, pathname, title, alternate, id, author):
        """
        Setup to manage an Atom feed.
        """
        self.pathname = pathname
        self.title = title
        self.alternate = alternate
        self.author = author
        try:
            self.etree = ElementTree.parse(self.pathname)
        except IOError:
            feed = ElementTree.Element(qname("feed"))
            ElementTree.SubElement(feed, qname("title"), {
                "type": "html"
            }).text = self.title
            ElementTree.SubElement(feed, qname("link"), {
                "href": id,
                "rel": "self"
            })
            ElementTree.SubElement(feed, qname("link"), {
                "href": alternate,
                "rel": "alternate"
            })
            ElementTree.SubElement(feed, "id").text = id
            ElementTree.SubElement(feed, qname("updated"))
            ElementTree.SubElement(ElementTree.SubElement(feed,
                qname("author")), qname("name")).text = self.author
            self.etree = ElementTree.ElementTree(feed)
            self.etree.write(self.pathname)
            self.etree = ElementTree.parse(self.pathname)

    def __del__(self):
        """
        Write the Atom feed back to disk with an updated timestamp.
        """
        self.etree.find(qname("updated")).text \
            = datetime.today().strftime("%Y-%m-%dT%H:%M:%SZ")
        self.etree.write(self.pathname)

    def entry(self,
        title=None,
        alternate=None,
        id=None,
        author=None,
        content=""
    ):
        """
        Create a new <entry>, bumping an old one out of the Atom feed if
        necessary.  Fill the new <entry> with the parameters or sensible
        defaults.
        """

        # Find sensible defaults for omitted arguments.
        if title is None:
            title = self.title
        if alternate is None:
            alternate = self.alternate
        if author is None:
            author = self.author

        # Setup the new <entry>.
        entry = self.etree.find(qname("entry"))
        if entry is None:
            entry = ElementTree.SubElement(self.etree.getroot(),
                qname("entry"))
        else:
            i = 0
            for e in self.etree.getroot().getchildren():
                if qname("entry") == e.tag:
                    break
                i += 1
            self.etree.getroot()[i+1:i+self.ENTRIES] \
                = self.etree.getroot()[i:i+self.ENTRIES-1]
            entry = self.etree.getroot()[i] = ElementTree.Element(
                qname("entry"))

        # Place the new article in the <entry>.
        ElementTree.SubElement(entry, qname("title"), {
            "type": "html"
        }).text = title
        ElementTree.SubElement(entry, qname("link"), {
            "href": alternate,
            "rel": "alternate"
        })
        ElementTree.SubElement(entry, qname("id")).text = id
        ElementTree.SubElement(entry, qname("published")).text \
            = datetime.today().strftime("%Y-%m-%dT%H:%M:%SZ")
        ElementTree.SubElement(entry, qname("updated")).text \
            = datetime.today().strftime("%Y-%m-%dT%H:%M:%SZ")
        ElementTree.SubElement(ElementTree.SubElement(entry,
            qname("author")), qname("name")).text = author
        ElementTree.SubElement(entry, qname("content"), {
            "type": "html"
        }).text = content


import pytest

from simplerestapi.api import GetUpdateDeleteAPI, ListCreateAPI
from simplerestapi.exception import ConfigEndpointError
from simplerestapi.main import SimpleApi
from tests import models

from .conftest import ERROR_TEMPLATE

class TestMain:
    @pytest.fixture(autouse=True)
    def setup(self, db_url_test):
        self.simple_api = SimpleApi(models, db_url_test)

    def test_get_models(self):
        assert len(self.simple_api.models) == 2
        assert models.Car in self.simple_api.models
        assert models.CustomUser in self.simple_api.models

    def test_construct_routes(self):
        expect_routes = [
            {'path': '/car', 'endpoint': ListCreateAPI, 'model': models.Car},
            {'path': '/car/{id}', 'endpoint': GetUpdateDeleteAPI, 'model': models.Car},
            {
                'path': '/customuser',
                'endpoint': ListCreateAPI,
                'model': models.CustomUser,
            },
            {
                'path': '/customuser/{id}',
                'endpoint': GetUpdateDeleteAPI,
                'model': models.CustomUser,
            },
        ]
        assert len(self.simple_api.routes) == 4
        result = [{'path': r.path, 'endpoint': r.endpoint, 'model': r.model} for r in self.simple_api.routes]
        for expect in expect_routes:
            assert expect in result

def test_raise_exceptions_by_errors(db_url_test):
    models.Car.ConfigEndpoint.pagination = 'not int'
    models.CustomUser.ConfigEndpoint.denied_methods = ['not_exist_method']
    with pytest.raises(ConfigEndpointError) as exc:
        SimpleApi(models, db_url_test)
    expected_errors = [
        ERROR_TEMPLATE.format('Car', 'pagination', 'value is not a valid integer'),
        ERROR_TEMPLATE.format(
            'CustomUser',
            'denied_methods',
            'methods must be from list (\'get\', \'post\', \'delete\', \'put\', \'patch\')',
        ),
    ]
    assert exc.value.args[0] == 'Your models have errors:\n' + '\n'.join(expected_errors)
    assert exc.value.errors == expected_errors


import random

from django.http import HttpResponse
from django.shortcuts import render, redirect
from django.views import View

class FormView(View):
    def get(self, request):
        return HttpResponse('Jesteśmy w metodzie GET')
    def post(selfself, request):
        return HttpResponse('Jesteśmy w metodzie POST')

class PizzaView(View):
    def get(self, request):
        return render(request, 'firstApp/pizza.html')
    def post(selfself, request):
        products = request.POST.getlist('skladnik', '')
        return HttpResponse(" ".join(products))
class CarView(View):
    def get(self, request):
        return render(request, 'firstApp/car.html')
    def post(selfself, request):
        car = request.POST.get('car', '')
        return HttpResponse(car)

class LoginView(View):
    def get(self, request):
        return render(request, 'firstApp/login.html')
    def post(selfself, request):
        login = request.POST.get('login', '')
        haslo = request.POST.get('haslo', '')
        if login == 'Admin' and haslo == 'Tajne123':
            return HttpResponse("Witaj admin!", 'brak parametru')
        else:
            return HttpResponse("Zły login lub hasło!", 'brak parametru')

def hello(request):
    return HttpResponse('Hello')
def adding(req):
    l1 = float(req.GET.get('l1', 0))
    l2 = float(req.GET.get('l2', 0))
    w = l1+l2
    return HttpResponse("Wynik: "+str(w))

def multiply(req):
    l = req.GET.get('l',0)
    wynik = "
"
    for i in range(1, int(l)+1):
        wynik += "
"
        for j in range(1,int(l)+1):
            wynik+=""
            wynik+=str(i*j)
            wynik+="	"
            j+=1
        wynik += "
"
        i+=1
    wynik += "
"
    return HttpResponse(wynik)

def brothers(req):
    names=[]
    wynik = ""
    names = req.GET.getlist('names','0')
    for i in names:
        wynik+="
"
        wynik+=i
        wynik+="
"
    return HttpResponse(wynik)

def fibonacci(req):
    l = int(req.GET.get('l', 0))
    w = [1,1]
    wynik = "
"
    wynik+="

1

1
"
    for i in range(2,l):
        temp = w[i-2]+w[i-1]
        wynik+="
"+str(temp)+"
"
        w.append(temp)
        i+=1
    wynik+="
"
    return HttpResponse(wynik)

def genFibo(n):
    a = 0
    b = 1
    for i in range(n):
        a , b = b , a + b
        yield a

def showSimpleFibo(req):
    l = int(req.GET.get('l', 0))
    wynik=""
    for i in genFibo(l):
        wynik+=str(i)
    return HttpResponse(wynik)

number = random.randint(1,100)
def game(req):
    global number
    message=""
    guess=-1
    try:
        guess = int(req.GET.get('guess', 0))
        if guess==0:
            odp = "<form><input name='guess' placeholder='podaj liczbę(1-100)'></form>"
        else:
            if(guess>number): message = "Twoja liczba jest za duża"
            elif(guess<number): message = "Twoja liczba jest za mała"
            else:
                message="zgadłeś/aś"
                number = random.randint(1, 100)
            odp = "<form><input name='guess' placeholder='podaj liczbę(1-100)'></form>" + message
        return HttpResponse(odp)
    except:
        pass

def helloPath(req,name):
    return HttpResponse(name)

def article(req, id):
    wynik=""
    if(id<4):
        wynik="Id to "+str(id)
    else:
        wynik= "brak artykułu o podanym id"
    return HttpResponse(wynik)

def greetings(req,name,repeat):
    wynik="Witaj "
    for i in range(0,repeat):
        wynik+=name
        wynik+=" "
        i+=1
    return HttpResponse(wynik)

def calc(req, nA, o, nB):
    try:
        if o == 'plus':
            return HttpResponse(str(nA+nB))
        elif o == 'minus':
            return HttpResponse(str(nA - nB))
        elif o == 'multiply':
            return HttpResponse(str(nA * nB))
        elif o == 'divide':
            if nB == 0:
                return HttpResponse('Nieprawidłowe dane')
            return HttpResponse(str(nA / nB))
    except:
        return HttpResponse('Nieprawidłowe dane')

def randomGenerator(req,min,max,throw=1):
    wynik=0
    while(throw):
        wynik+=random.randint(min, max)
        throw-=1
    return HttpResponse(str(wynik))

def index(request):
    return render(request,
                  'firstApp/index.html',
                  context={
          'msg':'wiadomość',
          'names':['Maja','Ala','Maria','Ewa'],
      })

def form(request):
    if request.method=='GET':
        return render(request, 'firstApp/form.html')
    if request.method == 'POST':
        return HttpResponse(request.POST.get("message",'brak parametru'))

def fizzBuzz(request):
    wynik=[]
    n = int(request.GET.get('n', 0))
    for i in range(1,n):
        if(i%3==0):
            if (i % 5 == 0):
                wynik.append("Fizz Buzz")
            else:
                wynik.append("Fizz")
        elif(i%5==0):
            wynik.append("Buzz")
        else:
            wynik.append(i)
    return render(request, 'firstApp/fizz_buzz.html',
                  context={'data':wynik})

def multiplyT(request, n):
    return render(request,'firstApp/multiply.html',context={'data':range(1,n+1)})

def rpg_game(request):
    player_hp = random.randint(100, 200)
    enemy_hp = random.randint(100, 500)

    stats=[]
    winner=[]
    player_hp_info = [player_hp]
    enemy_hp_info = [enemy_hp]

    for i in range(1,100):
        if enemy_hp<1 or player_hp<1:
            if enemy_hp<1: winner.append('player')
            elif player_hp<1: winner.append('enemy')
            break
        player_attack = random.randint(1, 100)
        enemy_hp -= player_attack

        enemy_attack = random.randint(10, 20)
        player_hp -= enemy_attack

        stats.append("Runda: "+str(i))
        stats.append("enemy lost: "+str(player_attack))
        stats.append("enemy life: "+str(enemy_hp))
        stats.append("player lost: " + str(enemy_attack))
        stats.append("player life: " + str(player_hp))

    return render(request, 'firstApp/rpg_game.html',
                  context={
                      'stats': stats,
                      'winner':winner,
                  })
def main(request):
    return render(request, 'firstApp/main.html')
def login_user(request, message="ok"):
    login = request.POST.get('login', '')
    haslo = request.POST.get('haslo', '')
    if request.method=='GET':
        return render(request, 'firstApp/login.html')
    if request.method == 'POST':
        if login=='Admin' and haslo=='Tajne123':
            return HttpResponse("Witaj admin!", 'brak parametru')
        else: return redirect('/login/?message="zły login lub hasło"')

products = {}

def add_product(request):
    product = request.POST.get('product', '')
    price = request.POST.get('price', '')
    if request.method=='GET':
        return render(request, 'firstApp/add_product_form.html')
    if request.method == 'POST':
        products[product] = price
        return redirect('/product/show')

def show_products(request):
    return render(request, 'firstApp/product_show.html',
                  context={'products':products})


from rest_framework import serializers

from goods.models import SKU

class CartSerializer(serializers.Serializer):
    sku_id = serializers.IntegerField(label="商品ID", min_value=1, required=True)
    count = serializers.IntegerField(label="购买数量", min_value=1)
    selected = serializers.BooleanField(label="时候勾选", default=True)

    def validate(self, attrs):
        """校验商品存不存在"""
        try:
            SKU.objects.get(id=attrs['sku_id'])
        except SKU.DoesNotExist:
            raise serializers.ValidationError('商品不存在')

        return attrs

class SKUCartSerializer(serializers.ModelSerializer):
    """购物车查询序列化器"""

    count = serializers.IntegerField(label='购买数量')
    selected = serializers.BooleanField(label='是否勾选')

    class Meta:
        model = SKU
        fields = ('id', 'count', 'name', 'default_image_url', 'price', 'selected')

class CartDeletedSerializer(serializers.Serializer):
    """购物车删除序列化器"""
    sku_id = serializers.IntegerField(label="商品ID", min_value=1, required=True)

    def validate(self, attrs):
        """校验商品存不存在"""
        try:
            SKU.objects.get(id=attrs['sku_id'])
        except SKU.DoesNotExist:
            raise serializers.ValidationError('商品不存在')

        return attrs


from annotation_utils.coco.structs import COCO_Dataset
from common_utils.path_utils import get_script_dir
import argparse

parser = argparse.ArgumentParser(description='Fix coco json file for pasonatron package')

parser.add_argument('img_dir', type=str,
                    help="path to image directory")
parser.add_argument('coco_ann_path', type=str,
                    help="path to json file")
parser.add_argument('coco_ann_output', type=str,
                    help="path to save a fixed json file")

args = parser.parse_args()

# img_dir = "/home/pasonatech/Blender2Detectron/BlenderProc-master/projectCrescent/crescent_test2/output3/coco_data/"
# coco_ann_path = "/home/pasonatech/Blender2Detectron/BlenderProc-master/projectCrescent/crescent_test2/output3/coco_data/coco_annotations.json"
# coco_ann_output ="/home/pasonatech/Blender2Detectron/BlenderProc-master/projectCrescent/crescent_test2/output3/coco_data/coco_annotations1.json"

# dataset = COCO_Dataset.load_from_path(coco_ann_path, strict=False, img_dir=get_script_dir())
dataset = COCO_Dataset.load_from_path(args.coco_ann_path, strict=False, img_dir=args.img_dir)
dataset.save_to_path(args.coco_ann_output, strict=True, overwrite=True)
print("Coco Json Fixed Successfully ")


from datetime import datetime
from decorators import debug_log

@debug_log
def validate_request(request):
    if 'action' in request and 'time' in request:
        return True
    return False

@debug_log
def make_response(request, code, data=None):
    return{
        'action': request.get('action'),
        'time': datetime.now().timestamp(),
        'code': code,
        'data': data
    }