import argparse
import datetime
import gym
import envs

import numpy as np
import torch
import imageio
import itertools

from rl.model import GaussianPolicy, QNetwork, DeterministicPolicy

from transformer_split.util import getGraphStructure
from transformer_split.vae_model import VAE_Model
from torch.nn import functional as F
from transformer_vae import util

parser = argparse.ArgumentParser(description='PyTorch Soft Actor-Critic Args')
parser.add_argument('--env1-name', default="ant",
                    help='Mujoco Gym environment (default: HalfCheetah-v2)')
parser.add_argument('--env2-name', default="ant3",
                    help='Mujoco Gym environment (default: HalfCheetah-v2)')
parser.add_argument('--model_path', default="runs/2021-05-19_13-46-41_VAE_ant-v0_both/",
                    help='model path')
parser.add_argument('--seed', type=int, default=123456, metavar='N',
                    help='random seed (default: 123456)')
parser.add_argument('--policy_hidden_size', type=int, default=256, metavar='N',
                    help='hidden size (default: 256)')
parser.add_argument('--latent_dim', type=int, default=128,
                    help='Encoder latent dimension')
parser.add_argument('--cuda', action="store_true",
                    help='run on CUDA (default: False)')
parser.add_argument('--agent_memory1', default='data/ant_jump.memory',
                    help='Path for saved replay memory')
parser.add_argument('--video_file_name', default="ant_turn.mp4",
                    help='output file name')
parser.add_argument('--msg_dim', type=int, default=32,
                        help='run on CUDA (default: False)')
parser.add_argument('--batch_size', type=int, default=1,
                        help='run on CUDA (default: False)')
parser.add_argument('--actor_path',
                    help='checkpoint training model every # steps')
parser.add_argument('--num_episodes', type=int, default=3, metavar='N',
                    help='maximum number of steps (default: 1000000)')
parser.add_argument('--root_size', type=int, default=11,
                        help='root dimension')
parser.add_argument('--lr', type=float, default=1e-4, metavar='N',
                        help='random seed (default: 123456)')
parser.add_argument(
    "--transformer_norm", default=0, type=int, help="Use layernorm",
)
parser.add_argument(
    "--beta",
    type=float,
    default=.1,
    help="beta coefficient of KL divergence",
)

parser.add_argument(
    "--gradient_penalty",
    type=float,
    default=10,
    help="beta coefficient of KL divergence",
)

parser.add_argument(
    "--discriminator_limiting_accuracy",
    type=float,
    default=0.7,
    help="beta coefficient of KL divergence",
)
parser.add_argument(
    "--attention_layers",
    default=3,
    type=int,
    help="How many attention layers to stack",
)
parser.add_argument(
    "--attention_heads",
    default=2,
    type=int,
    help="How many attention heads to stack",
)
parser.add_argument(
    "--attention_hidden_size",
    type=int,
    default=128,
    help="Hidden units in an attention block",
)

parser.add_argument(
    "--attention_embedding_size",
    type=int,
    default=128,
    help="Hidden units in an attention block",
)

parser.add_argument(
    "--dropout_rate",
    type=float,
    default=0.0,
    help="How much to drop if drop in transformers",
)
args = parser.parse_args()

torch.manual_seed(args.seed)
np.random.seed(args.seed)

# Agent
device = torch.device("cuda" if args.cuda else "cpu")

env_names = ["ant-v0", "ant3-v0", "ant_a-v0"]
train_envs = [gym.make(n) for n in env_names]
graphs = [getGraphStructure(e.xml) for e in train_envs]
# All environments have the same dimension per limb.
num_limbs = len(graphs[0])  #torso + body limbs
body_limbs = num_limbs - 1
dim_per_limb = int((train_envs[0].observation_space.shape[0] - args.root_size) / (body_limbs - 1))
max_num_limbs = max(len(g) for g in graphs)

args.dim_per_limb = dim_per_limb
args.max_num_limbs = max_num_limbs

root_dir = util.get_project_root()
render_env = train_envs[2]
render_topology = graphs[2]
render_limbs = len(render_topology)

expert_env = train_envs[0]
expert_topology = graphs[0]

policy = GaussianPolicy(
            expert_env.observation_space.shape[0],
            expert_env.action_space.shape[0],
            args.policy_hidden_size,
            expert_env.action_space).to(device)

policy.load_state_dict(torch.load(args.actor_path))

vae_model = VAE_Model(args)

vae_model.load_model(args.model_path)

def pad_state(data, state_size, max_num_limbs):
    max_dim = args.root_size + state_size * (max_num_limbs - 1)
    output = torch.zeros(max_dim)
    output[:data.shape[0]] = torch.tensor(data)
    return output

def pad_topology(top, max_num_limbs):
    topology = torch.full((max_num_limbs,), -1, dtype=torch.int32)
    topology[:len(top)] = torch.tensor(top, dtype=torch.int32)
    return topology

# Evaluation loop
total_numsteps = 0
avg_reward = 0.
state = render_env.reset()

with imageio.get_writer(args.video_file_name, fps=30) as video:

    for i_episode in itertools.count(1):
        episode_reward = 0
        episode_steps = 0
        done = False

        state = render_env.reset()
        video.append_data(render_env.render('rgb_array'))

        done = False
        while not done:
            state = pad_state(state, dim_per_limb, max_num_limbs).unsqueeze(0)
            src_topology = pad_topology(render_topology, max_num_limbs).unsqueeze(0)
            tgt_topology = pad_topology(expert_topology, max_num_limbs).unsqueeze(0)

            x_hat = vae_model.transfer(state, tgt_topology)
            x_hat = x_hat.detach().cpu()
            x_hat = x_hat[:(render_limbs-1)]
            x_hat = torch.FloatTensor(x_hat).to(device).unsqueeze(0)
            action, _, _ = policy.sample(x_hat)
            action = action.detach().cpu().numpy()[0]

            next_state, reward, done, _ = render_env.step(action[0][:7])
            video.append_data(render_env.render('rgb_array'))

            episode_reward += reward

            state = next_state
        avg_reward += episode_reward

        if i_episode > args.num_episodes:
            break


from __future__ import division
import itertools
import logging
logger = logging.getLogger(__name__)

import numpy as np

from .structure import Ligand, BondOrder
from .volume import Volume
from .samplers import GlobalRotator, Translator, BondRotator, RotationSets, ClashDetector
from .transformer import Transformer
from .solvers import QPSolver, MIQPSolver
from .helpers import DJoiner

class HierarchicalBuilder(object):

    """Build a multi-conformer ligand hierarchically."""

    def __init__(self, ligand, xmap, resolution, receptor=None,
            local_search=True, build=True,
            stepsize=2, build_stepsize=1,
            threshold=None, cardinality=5,
            directory='.', roots=None, debug=False):
        self.ligand = ligand
        self.xmap = xmap
        self.resolution = resolution
        self.build = build
        self.stepsize= stepsize
        self.build_stepsize = build_stepsize
        self.directory = directory
        self.local_search = local_search
        self.receptor = receptor
        if self.resolution < 3.0:
            self._rmask = 0.5 + self.resolution / 3.0
        else:
            self._rmask = 0.5 * self.resolution
        self._debug = debug

        # For MIQP
        self.threshold = threshold
        self.cardinality = cardinality

        self._trans_box = [(-0.2, 0.21, 0.1)] * 3
        self._sampling_range = np.deg2rad(np.arange(0, 360, self.stepsize))
        self._djoiner = DJoiner(directory)

        if self.receptor is not None:
            self._cd = ClashDetector(self.ligand, self.receptor, scaling_factor=0.75)
            if self._cd():
                logger.warning("Initial ligand configuration is clashing!")

        self._rigid_clusters = self.ligand.rigid_clusters()
        # Determine which roots to start building from
        if roots is None:
            self._clusters_to_sample = []
            for cluster in self._rigid_clusters:
                nhydrogen = (self.ligand.e[cluster] == 'H').sum()
                if len(cluster) - nhydrogen > 1:
                    self._clusters_to_sample.append(cluster)
        else:
            self._clusters_to_sample = []
            for root in roots:
                for cluster in self._rigid_clusters:
                    if root in cluster:
                        self._clusters_to_sample.append(cluster)
        msg = "Number of clusters to sample: {:}".format(
                len(self._clusters_to_sample))
        logger.info(msg)
        self._starting_coor_set = [ligand.coor.copy()]
        self._coor_set = []
        self._all_coor_set = []
        self.conformers = []
        self._occupancies = []

        smax = 1 / (2 * self.resolution)
        self._model_map = Volume.zeros_like(self.xmap)
        # Initialize density creation
        # We can let go of the spacegroup now that we have prepared the map.
        self._model_map.set_spacegroup("P1")
        self._transformer = Transformer(self.ligand, self._model_map, smax=smax, rmax=3)

    def __call__(self):

        self._all_occupancies = []
        if self.build:
            for self._cluster_index, self._cluster in enumerate(self._clusters_to_sample):
                self._iteration = 0
                self._coor_set = list(self._starting_coor_set)
                logger.info("Cluster index: {:}".format(self._cluster_index))
                logger.info("Number of conformers: {:}".format(len(self._coor_set)))
                if self.local_search:
                    self._local_search()
                self._build_ligand()

                self._all_coor_set += self._coor_set
                self._all_occupancies += list(self._occupancies)
                logger.info("Number of conformers: {:}".format(len(self._coor_set)))
                logger.info("Number of final conformers: {:}".format(len(self._all_coor_set)))

        self._coor_set = self._all_coor_set
        self._occupancies = np.asarray(self._all_occupancies)
        self._convert()

    def _clashing(self):
        if self.receptor is None:
            return self.ligand.clashes()
        else:
            return self.ligand.clashes() or self._cd()

    def _local_search(self):
        """Perform a local rigid body search on the cluster."""

        logger.info("Performing local search.")

        # Set occupancies of rigid cluster and its direct neighboring atoms to
        # 1 for clash detection and MIQP
        self.ligand.q.fill(0)
        self.ligand.q[self._cluster] = 1
        for atom in self._cluster:
            self.ligand.q[self.ligand.connectivity[atom]] = 1
        center = self.ligand.coor[self._cluster].mean(axis=0)
        new_coor_set = []
        for coor in self._coor_set:
            self.ligand.coor[:] = coor
            rotator = GlobalRotator(self.ligand, center=center)
            for rotmat in RotationSets.get_local_set():
                rotator(rotmat)
                translator = Translator(self.ligand)
                iterator = itertools.product(*[
                    np.arange(*trans) for trans in self._trans_box])
                for translation in iterator:
                        translator(translation)
                        if not self._clashing():
                            new_coor_set.append(self.ligand.coor.copy())
        self._coor_set = new_coor_set
        # In case all conformers were clashing
        if not self._coor_set:
            return
        if self._debug:
            self._write_intermediate_structures(base='intermediate')
        self._convert()
        self._QP()
        self._update_conformers()
        if self._debug:
            self._write_intermediate_structures(base='qp')
        self._convert()
        self._MIQP()
        self._update_conformers()
        if self._debug:
            self._write_intermediate_structures(base='miqp')

    def _build_ligand(self):
        """Build up the ligand hierarchically."""

        logger.info("Building up ligand.")
        # Sampling order of bonds
        bond_order = BondOrder(self.ligand, self._cluster[0])
        bonds = bond_order.order
        depths = bond_order.depth
        nbonds = len(bonds)
        logger.info("Number of bonds to sample: {:d}".format(nbonds))
        starting_bond_index = 0
        finished_building = True if nbonds == 0 else False
        while not finished_building:
            end_bond_index = min(starting_bond_index + self.build_stepsize, nbonds)
            logger.info("Sampling iteration: {:d}".format(self._iteration))
            for bond_index in xrange(starting_bond_index, end_bond_index):
                # Set the occupancies of build clusters and their direct
                # neighbors to 1 for clash detection and MIQP.
                nbonds_sampled = bond_index + 1
                self.ligand.q.fill(0)
                for cluster in self._rigid_clusters:
                    for sampled_bond in bonds[:nbonds_sampled]:
                        if sampled_bond[0] in cluster or sampled_bond[1] in cluster:
                            self.ligand.q[cluster] = 1
                            for atom in cluster:
                                self.ligand.q[self.ligand.connectivity[atom]] = 1

                # Sample by rotating around a bond
                bond = bonds[bond_index]
                atoms = [self.ligand.atomname[bond[0]], self.ligand.atomname[bond[1]]]
                new_coor_set = []
                for coor in self._coor_set:
                    self.ligand.coor[:] = coor
                    rotator = BondRotator(self.ligand, *atoms)
                    for angle in self._sampling_range:
                        rotator(angle)
                        if not self._clashing():
                            new_coor_set.append(self.ligand.coor.copy())
                self._coor_set = new_coor_set
                # Check if any acceptable configurations have been created.
                if not self._coor_set:
                    msg = "No non-clashing conformers found. Stop building."
                    logger.warning(msg)
                    return

                # Perform an MIQP if either the end bond index has been reached
                # or if the end of a sidechain has been reached, i.e. if the
                # next depth level is equal or smaller than current depth. If
                # we step out of bonds in the depths list it means we are done.
                end_iteration = (nbonds_sampled == end_bond_index)
                try:
                    end_sidechain = depths[bond_index] >= depths[bond_index + 1]
                except IndexError:
                    finished_building = True
                    end_sidechain = True
                if end_iteration or end_sidechain:
                    self._iteration += 1
                    if self._debug:
                        self._write_intermediate_structures()
                    self._convert()
                    self._QP()
                    self._update_conformers()
                    if self._debug:
                        self._write_intermediate_structures('qp')
                    self._convert()
                    self._MIQP()
                    self._update_conformers()
                    if self._debug:
                        self._write_intermediate_structures('miqp')

                    # Stop this building iteration and move on to next
                    starting_bond_index += 1
                    if end_sidechain:
                        starting_bond_index = nbonds_sampled
                    break

    def _convert(self, tofile=False):

        logger.info('Converting structures to densities ({:})'.format(len(self._coor_set)))
        self._transformer.volume.array.fill(0)
        for coor in self._coor_set:
            self.ligand.coor[:] = coor
            self._transformer.mask(self._rmask)
        if tofile:
            self._transformer.volume.tofile('mask.ccp4')
        mask = self._transformer.volume.array > 0
        self._transformer.volume.array.fill(0)

        nvalues = mask.sum()
        self._target = self.xmap.array[mask]
        self._models = np.zeros((len(self._coor_set), nvalues), dtype=np.float64)
        for n, coor in enumerate(self._coor_set):
            self.ligand.coor[:] = coor
            self._transformer.density()
            self._models[n] = self._transformer.volume.array[mask]
            if tofile:
                self._transformer.volume.tofile('density_{}.ccp4'.format(n))
            self._transformer.reset()

    def _QP(self):
        logger.info("Starting QP.")
        qpsolver = QPSolver(self._target, self._models)
        logger.info("Initializing.")
        qpsolver.initialize()
        logger.info("Solving")
        qpsolver()
        self._occupancies = qpsolver.occupancies

    def _MIQP(self, maxfits=None, exact=False, threshold=None):
        logger.info("Starting MIQP.")
        miqpsolver = MIQPSolver(self._target, self._models)
        logger.info("Initializing.")
        miqpsolver.initialize()
        logger.info("Solving")
        if maxfits is None:
            cardinality = self.cardinality
        else:
            cardinality = maxfits
        if threshold is None:
            threshold = self.threshold
        else:
            threshold = threshold
        miqpsolver(maxfits=cardinality, exact=exact, threshold=threshold)
        self._occupancies = miqpsolver.occupancies

    def _update_conformers(self):
        logger.info("Updating conformer list.")
        logger.info("Old number of conformers: {:d}".format(len(self._coor_set)))
        new_coor_set = []
        cutoff = 0.002
        for n, coor in enumerate(self._coor_set):
            if self._occupancies[n] >= cutoff:
                new_coor_set.append(coor)
        if new_coor_set:
            self._coor_set = new_coor_set
            self._occupancies = self._occupancies[self._occupancies >= cutoff]
        else:
            logger.warning("No conformer found with occupancy bigger than {:}.".format(cutoff))
            sorted_indices = np.argsort(self._occupancies)[::-1]
            indices = np.arange(len(self._occupancies))[sorted_indices]
            nbest = min(5, len(self._occupancies))
            for i in indices[:nbest]:
                new_coor_set.append(self._coor_set[i])
            self._coor_set = new_coor_set
            self._occupancies = self._occupancies[indices[:nbest]]
        logger.info("New number of conformers: {:d}".format(len(self._coor_set)))

    def _write_intermediate_structures(self, base='intermediate'):
        logger.info("Writing intermediate structures to file.")
        fname_base = self._djoiner(base + '_{:d}_{:d}_{:d}.pdb')
        for n, coor in enumerate(self._coor_set):
            self.ligand.coor[:] = coor
            ligand = self.ligand.select('q', 0, '!=')
            fname = fname_base.format(self._cluster_index, self._iteration, n)
            ligand.tofile(fname)

    def get_conformers(self, cutoff=0.01):
        conformers = []
        iterator = itertools.izip(self._coor_set, self._occupancies)
        for coor, occ in iterator:
            if occ >= cutoff:
                ligand = Ligand(self.ligand.data.copy(), self.ligand.coor.copy())
                ligand.coor[:] = coor
                ligand.q.fill(occ)
                conformers.append(ligand)
        ## Sort conformers based on occupancy
        #conformers = sorted(conformers, key=lambda conformer: conformer.q[0], reverse=True)
        return conformers

    def write_results(self, base='conformer', cutoff=0.01):
        logger.info("Writing results to file.")
        fname_base = self._djoiner(base + '_{:d}.pdb')
        fnames = []
        iterator = itertools.izip(self._coor_set, self._occupancies)
        n = 1
        old_q = self.ligand.q.copy()
        for coor, occ in iterator:
            if occ >= cutoff:
                self.ligand.q.fill(occ)
                self.ligand.coor[:] = coor
                fname = fname_base.format(n)
                self.ligand.tofile(fname)
                n += 1
                fnames.append(fname)
        self.ligand.q[:] = old_q
        return fnames


import audioop
import numpy as np
import pyaudio
import wave

CHUNK_SIZE = 1024

class Music:

    def __init__(self, path = "./stereo.wav"):
        self.path = path
        self.wf = wave.open(path, 'rb')
        self.width = self.wf.getsampwidth()
        self.pa = pyaudio.PyAudio()
        self.stream = self.pa.open(format = self.pa.get_format_from_width(self.width),
                                   channels = 6,
                                   rate = self.wf.getframerate(),
                                   output = True)

    def stop(self):
        self.stream.stop_stream()
        self.stream.close()
        self.pa.terminate()

    def replay(self):
        # バイト列を取得
        # [L0, R0, L1, R1, L2, R2, ...]
        src_frames = self.wf.readframes(CHUNK_SIZE)

        while src_frames != '':
            # バイト列を取得
            # [L0, R0, L1, R1, L2, R2, ...]
            src_frames = self.wf.readframes(CHUNK_SIZE)

            # L, Rに分割
            l_frames = audioop.tomono(src_frames, self.width, 1, 0)
            r_frames = audioop.tomono(src_frames, self.width, 0, 1)

            # 各chに流す音量を調整
            # L, FL, FR, R, Bの順
            volumes = [[1, 1, 1, 1, 1], [0, 0, 0, 0, 0]]
            six_ch_frames = self.set_6ch_audio(l_frames, r_frames, volumes)

            # 6chオーディオをstreamに渡す
            # [FL0, FR0, CT0, BA0, RL0, RR0, ...]
            self.stream.write(six_ch_frames)

        self.stop()

    def set_6ch_audio(self, l_frames, r_frames, volumes):
        all_frame = [ [np.fromstring(audioop.mul(l_frames, self.width, volumes[0][i]), dtype=np.int16)]
                      if volumes[0][i] != 0
                      else [np.fromstring(audioop.mul(r_frames, self.width, volumes[1][i]), dtype=np.int16)]
                      for i in range(5)]
        audio = np.concatenate(all_frame, axis=0)
        # insertの第2引数がミュートするチャンネル
        audio = np.insert(audio, 4, 0, axis=0).T.astype(np.int16).tostring()

        return audio

if __name__ == '__main__':
    music = Music()
    music.replay()


import logging
from collections import defaultdict
from typing import Any, Callable, Dict, Iterator, List, NamedTuple, Optional, Tuple, Union

import numpy as np
import torch
from ptgnn.baseneuralmodel import AbstractNeuralModel, ModuleWithMetrics
from ptgnn.neuralmodels.embeddings.strelementrepresentationmodel import (
    CharUnitEmbedder,
    StrElementRepresentationModel,
    SubtokenUnitEmbedder,
    TokenUnitEmbedder,
)
from torch import nn
from typing_extensions import Literal

from buglab.models.basemodel import AbstractBugLabModel
from buglab.models.layers.fixermodules import (
    CandidatePairSelectorModule,
    SingleCandidateNodeSelectorModule,
    TextRepairModule,
)
from buglab.models.layers.localizationmodule import LocalizationModule
from buglab.models.layers.relational_transformer import RelationalTransformerEncoderLayer
from buglab.models.utils import compute_generator_loss, scatter_log_softmax
from buglab.representations.data import BugLabData, BugLabGraph

LOGGER = logging.getLogger(__name__)

class SeqModelTensorizedSample(NamedTuple):
    target_subtokens_ids: List[int]
    intra_token_edges: Dict[str, List[Tuple[int, int]]]
    candidate_location_idxs: List[int]

    target_location_idx: Optional[int]

    node_mappings: Dict[int, int]

    # Bug repair
    target_rewrite_node_ids: List[int]
    target_rewrites: List[int]
    target_rewrite_to_location_group: List[int]
    correct_rewrite_target: Optional[int]
    text_rewrite_original_idx: List[int]

    varmisused_node_ids: List[int]
    candidate_symbol_node_ids: List[int]
    candidate_symbol_to_varmisused_node: List[int]
    correct_candidate_symbol_node: Optional[int]
    candidate_rewrite_original_idx: List[int]

    call_node_ids: List[int]
    candidate_swapped_node_ids: List[int]
    swapped_pair_to_call: List[int]
    correct_swapped_pair: Optional[int]
    pair_rewrite_original_idx: List[int]

    num_rewrite_locations_considered: int

    # Bug selection
    rewrite_logprobs: Optional[List[float]]

class SeqBugLabModule(ModuleWithMetrics):
    def __init__(
        self,
        token_embedder: Union[TokenUnitEmbedder, SubtokenUnitEmbedder, CharUnitEmbedder],
        embedding_dim: int,
        num_edge_types: int,
        num_layers: int,
        num_heads: int,
        intermediate_dimension: int,
        dropout_rate: float,
        rewrite_vocabulary_size: int,
        layer_type: Literal["great", "rat", "transformer", "gru"] = "great",
        buggy_samples_weight_schedule: Callable[[int], float] = lambda _: 1.0,
        generator_loss_type: Optional[str] = "norm-kl",
        rezero_mode: Literal["off", "scalar", "vector"] = "off",
        normalisation_mode: Literal["off", "prenorm", "postnorm"] = "postnorm",
    ):
        super().__init__()
        self.__generator_loss_type = generator_loss_type
        self.__token_embedder = token_embedder

        self.__positional_encoding = nn.Parameter(torch.randn(1, 5000, embedding_dim), requires_grad=True)

        self.__dropout_layer = nn.Dropout(dropout_rate)
        self.__input_layer_norm = nn.LayerNorm(embedding_dim)

        self.__layer_type = layer_type
        if layer_type in {"great", "rat"}:
            self.__seq_layers = nn.ModuleList(
                [
                    RelationalTransformerEncoderLayer(
                        nhead=num_heads,
                        num_edge_types=num_edge_types,
                        d_model=embedding_dim,
                        key_query_dimension=embedding_dim // num_heads,
                        value_dimension=embedding_dim // num_heads,
                        dim_feedforward=intermediate_dimension,
                        dropout=dropout_rate,
                        use_edge_value_biases=layer_type == "rat",
                        rezero_mode=rezero_mode,
                        normalisation_mode=normalisation_mode,
                    )
                    for _ in range(num_layers)
                ]
            )
        elif layer_type == "transformer":
            self.__seq_layers = nn.ModuleList(
                [
                    nn.TransformerEncoderLayer(
                        d_model=embedding_dim,
                        nhead=num_heads,
                        dim_feedforward=intermediate_dimension,
                        dropout=dropout_rate,
                    )
                    for _ in range(num_layers)
                ]
            )
        elif layer_type == "gru":
            self.__seq_layers = nn.GRU(
                input_size=embedding_dim,
                hidden_size=embedding_dim // 2,
                num_layers=num_layers,
                bidirectional=True,
                batch_first=True,
            )
        else:
            raise ValueError(f"Unrecognized layer type `{layer_type}`.")

        output_representation_size = embedding_dim
        self.__localization_module = LocalizationModule(
            output_representation_size, buggy_samples_weight_schedule=buggy_samples_weight_schedule
        )
        self._buggy_samples_weight_schedule = buggy_samples_weight_schedule

        self._text_repair_module = TextRepairModule(embedding_dim, rewrite_vocabulary_size)
        self._varmisuse_module = SingleCandidateNodeSelectorModule(embedding_dim)
        self._argswap_module = CandidatePairSelectorModule(embedding_dim)

    def _reset_module_metrics(self) -> None:
        if not hasattr(self, "_epoch_idx"):
            self._epoch_idx = 0
        elif self.training and self.__num_batches > 0:
            # Assumes that module metrics are reset once per epoch.
            self._epoch_idx += 1

        self.__loss = 0.0
        self.__repair_loss = 0.0

        self.__total_samples = 0
        self.__num_batches = 0

    def _module_metrics(self) -> Dict[str, Any]:
        metrics = {}
        if self.__total_samples > 0:
            metrics["Repair Loss"] = self.__repair_loss / self.__total_samples
        if self.__num_batches > 0:
            metrics["Loss"] = self.__loss / self.__num_batches
        return metrics

    def _compute_repair_logprobs(
        self,
        output_representations,
        target_rewrite_node_ids,
        target_rewrites,
        rewrite_to_location_group,
        varmisused_node_ids,
        candidate_symbol_node_ids,
        candidate_symbol_to_location_group,
        call_node_ids,
        candidate_swapped_node_ids,
        swapped_pair_to_call_location_group,
    ):
        if target_rewrites.shape[0] > 0:
            text_repair_logits = self._text_repair_module.compute_rewrite_logits(
                target_rewrite_node_representations=output_representations[
                    target_rewrite_node_ids[:, 0], target_rewrite_node_ids[:, 1]
                ],
                candidate_rewrites=target_rewrites,
            )
        else:
            text_repair_logits = torch.zeros(0, device=target_rewrites.device)

        if varmisused_node_ids.shape[0] > 0:
            varmisuse_logits = self._varmisuse_module.compute_per_slot_log_probability(
                slot_representations_per_target=output_representations[
                    varmisused_node_ids[:, 0], varmisused_node_ids[:, 1]
                ],
                target_nodes_representations=output_representations[
                    candidate_symbol_node_ids[:, 0], candidate_symbol_node_ids[:, 1]
                ],
            )
        else:
            varmisuse_logits = torch.zeros(0, device=varmisused_node_ids.device)

        if call_node_ids.shape[0] > 0:
            arg_swap_logits = self._argswap_module.compute_per_pair_logits(
                slot_representations_per_pair=output_representations[call_node_ids[:, 0], call_node_ids[:, 1]],
                pair_representations=output_representations[
                    candidate_swapped_node_ids[:, 0].unsqueeze(-1), candidate_swapped_node_ids[:, 1:]
                ],
            )
        else:
            arg_swap_logits = torch.zeros(0, device=call_node_ids.device)

        all_logits = torch.cat((text_repair_logits, varmisuse_logits, arg_swap_logits))
        logit_groups = torch.cat(
            (rewrite_to_location_group, candidate_symbol_to_location_group, swapped_pair_to_call_location_group)
        )
        logprobs = scatter_log_softmax(all_logits, index=logit_groups)
        text_repair_logprobs, varmisuse_logprobs, arg_swap_logprobs = torch.split(
            logprobs, [text_repair_logits.shape[0], varmisuse_logits.shape[0], arg_swap_logits.shape[0]]
        )
        return arg_swap_logprobs, text_repair_logprobs, varmisuse_logprobs

    def _compute_localization_logprobs(self, candidate_reprs, candidate_to_sample_idx, num_samples):
        candidate_to_sample_idx, candidate_log_probs, _ = self.__localization_module.compute_localization_logprobs(
            candidate_reprs, candidate_to_sample_idx, num_samples
        )
        return candidate_to_sample_idx, candidate_log_probs

    def forward(
        self,
        *,
        input_sequence_ids,
        input_seq_num_subtokens,
        token_sequence_lengths,
        edges,
        edge_types,
        has_bug,
        candidate_location_idxs,
        target_location_idxs,
        # Text repair
        target_rewrite_node_ids: torch.IntTensor,
        target_rewrites: torch.IntTensor,
        rewrite_to_location_group: torch.IntTensor,
        correct_rewrite_idxs: torch.IntTensor,
        text_rewrite_idxs: torch.IntTensor,
        # Var Misuse
        varmisused_node_ids: torch.IntTensor,
        candidate_symbol_node_ids: torch.IntTensor,
        candidate_symbol_to_location_group: torch.IntTensor,
        correct_candidate_symbols: torch.IntTensor,
        candidate_rewrite_idxs: torch.IntTensor,
        # ArgSwap
        call_node_ids: torch.IntTensor,
        candidate_swapped_node_ids: torch.IntTensor,
        swapped_pair_to_call_location_group: torch.IntTensor,
        correct_swapped_pair: torch.IntTensor,
        pair_rewrite_idxs: torch.IntTensor,
        # Rewrite logprobs
        rewrite_to_graph_id: Optional[torch.IntTensor] = None,
        rewrite_logprobs: Optional[torch.FloatTensor] = None,
        **kwargs,  # Visualization data (ignored)
    ):
        """
        :param input_sequence_ids: [B, max_seq_len]
        :param input_sequence_lengths: The size of the input sequence [B]
        :param token_sequence_lengths: The size of the tokens in the input sequence (excluding symbols) [B]
        :param edges: [num_edges, 3]
        :param edge_types: [num_edges]
        :param has_bug: [B]
        :param candidate_location_idxs: [num_target_locations, 2]
        :param target_location_idxs: [B]
        """
        output_representation = self._compute_output_representation(
            input_sequence_ids, input_seq_num_subtokens, token_sequence_lengths, edges, edge_types
        )

        target_location_representations = output_representation[
            candidate_location_idxs[:, 0], candidate_location_idxs[:, 1]
        ]  # [num_candidate_points, H]

        # Repair
        arg_swap_logprobs, text_repair_logprobs, varmisuse_logprobs = self._compute_repair_logprobs(
            output_representation,
            target_rewrite_node_ids,
            target_rewrites,
            rewrite_to_location_group,
            varmisused_node_ids,
            candidate_symbol_node_ids,
            candidate_symbol_to_location_group,
            call_node_ids,
            candidate_swapped_node_ids,
            swapped_pair_to_call_location_group,
        )

        if rewrite_logprobs is not None:
            candidate_to_sample_idx = candidate_location_idxs[:, 0]
            loss_type = self.__generator_loss_type
            # Note: the command
            # torch.exp(torch_scatter.scatter_logsumexp(localization_logprobs, torch.cat((candidate_to_sample_idx, arange))))
            # should give a tensor of ones here (total probabilities).
            _, localization_logprobs, arrange, = self.__localization_module.compute_localization_logprobs(
                candidate_reprs=target_location_representations,
                candidate_to_sample_idx=candidate_to_sample_idx,
                num_samples=has_bug.shape[0],
            )

            loss = compute_generator_loss(
                arg_swap_logprobs,
                arrange,
                candidate_rewrite_idxs,
                candidate_symbol_to_location_group,
                localization_logprobs,
                loss_type,
                pair_rewrite_idxs,
                rewrite_logprobs,
                rewrite_to_graph_id,
                rewrite_to_location_group,
                swapped_pair_to_call_location_group,
                text_repair_logprobs,
                text_rewrite_idxs,
                varmisuse_logprobs,
            )

            with torch.no_grad():
                self.__loss += float(loss)
                self.__num_batches += 1

            return loss

        else:
            localization_loss = self.__localization_module(
                candidate_reprs=target_location_representations,
                candidate_to_sample_idx=candidate_location_idxs[:, 0],
                has_bug=has_bug,
                correct_candidate_idxs=target_location_idxs,
            )

            text_repair_loss = self._text_repair_module(text_repair_logprobs, correct_rewrite_idxs)
            varmisuse_repair_loss = self._varmisuse_module(varmisuse_logprobs, correct_candidate_symbols)
            arg_swap_repair_loss = self._argswap_module(arg_swap_logprobs, correct_swapped_pair)

            buggy_samples_weight = self._buggy_samples_weight_schedule(self._epoch_idx)
            repair_loss = text_repair_loss.sum() + varmisuse_repair_loss.sum() + arg_swap_repair_loss.sum()
            repair_loss = repair_loss * buggy_samples_weight

            with torch.no_grad():
                self.__total_samples += float(has_bug.sum().cpu())
                self.__repair_loss += float(repair_loss)

                self.__num_batches += 1
                self.__loss += float(localization_loss + repair_loss / has_bug.shape[0])

            return localization_loss + repair_loss / has_bug.shape[0]

    def _compute_output_representation(
        self, input_sequence_ids, input_seq_num_subtokens, token_sequence_lengths, edges, edge_types
    ):
        output_representation = self.__token_embedder(
            token_idxs=input_sequence_ids.reshape(-1, input_sequence_ids.shape[-1]),
            lengths=input_seq_num_subtokens.reshape(-1),
        )  # [B*max_seq_len, D]

        output_representation = output_representation.reshape(
            input_sequence_ids.shape[0], input_sequence_ids.shape[1], -1
        )  # [B, max_seq_len, D]
        token_range = torch.arange(input_sequence_ids.shape[1], device=token_sequence_lengths.device)
        token_mask = token_range.unsqueeze(0) < token_sequence_lengths.unsqueeze(-1)  # [B, max_seq_len]

        if self.__layer_type in {"rat", "transformer", "great"}:
            # Add positional embeddings only to token nodes (and not to symbol nodes)

            output_representation = output_representation + self.__positional_encoding[:, : input_sequence_ids.shape[1]]

            # The original BERT model has this layer after embedding + positional encoding
            output_representation = self.__dropout_layer(self.__input_layer_norm(output_representation))
            seq_mask = token_mask.logical_not()  # [B, max_seq_len]

        output_representation *= token_mask.unsqueeze(-1)

        if self.__layer_type in {"rat", "great"}:
            for layer in self.__seq_layers:
                output_representation = layer(
                    src=output_representation, src_mask=seq_mask, edges=edges, edge_types=edge_types
                )  # [B, max_seq_len, H]
        elif self.__layer_type == "transformer":
            for layer in self.__seq_layers:
                output_representation = output_representation.transpose(0, 1)  # [max_seq_len, B, D]
                output_representation = layer(output_representation, src_key_padding_mask=seq_mask)
                output_representation = output_representation.transpose(0, 1)  # [B, max_seq_len, D]
        elif self.__layer_type == "gru":
            output_representation, _ = self.__seq_layers(
                nn.utils.rnn.pack_padded_sequence(
                    output_representation, lengths=token_sequence_lengths, batch_first=True, enforce_sorted=False
                )
            )
            output_representation, lens = nn.utils.rnn.pad_packed_sequence(output_representation, batch_first=True)
        else:
            raise ValueError()

        return output_representation

class SeqBugLabModel(AbstractNeuralModel[BugLabData, SeqModelTensorizedSample, SeqBugLabModule], AbstractBugLabModel):
    def __init__(
        self,
        representation_size: int,
        max_subtoken_vocab_size: int,
        dropout_rate: float,
        layer_type: Literal["great", "rat", "transformer", "gru"] = "great",
        max_seq_size: int = 500,
        num_heads: int = 8,
        num_layers: int = 6,
        intermediate_dimension_size: int = 2048,
        buggy_samples_weight_schedule: Callable[[int], float] = lambda _: 1.0,
        generator_loss_type: Optional[str] = "classify-max-loss",
        rezero_mode: Literal["off", "scalar", "vector"] = "off",
        normalisation_mode: Literal["off", "prenorm", "postnorm"] = "postnorm",
    ):
        super().__init__()
        self._init()
        self.__tensorize_only_at_target_location_rewrites = True
        self.__edge_types = set()

        self.__dropout_rate = dropout_rate
        self.__representation_size = representation_size
        self.__layer_type = layer_type

        self.__max_seq_size = max_seq_size
        self.__token_embedder = StrElementRepresentationModel(
            token_splitting="subtoken",
            embedding_size=representation_size,
            dropout_rate=dropout_rate,
            vocabulary_size=max_subtoken_vocab_size,
            subtoken_combination="max",
        )

        self.__buggy_samples_weight_schedule = buggy_samples_weight_schedule

        self.__num_heads = num_heads
        self.__num_layers = num_layers
        self.__intermediate_dimension_size = intermediate_dimension_size
        self.__generator_loss_type = generator_loss_type
        self.__rezero_mode = rezero_mode
        self.__normalisation_mode = normalisation_mode

    def __to_token_data(self, graph: BugLabGraph):
        # Get token list
        token_sequence, next_token_edges = self.__extract_token_sequence(graph)
        if len(token_sequence) == 0:
            return None
        old_node_id_to_new = {t: i for i, t in enumerate(token_sequence)}

        # Get all edges among tokens
        NON_TOKEN_EDGES = frozenset(
            {
                "NextToken",
                "PossibleType",
                "CandidateCall",
                "CandidateCallDoc",
                "MayFormalName",
                "Child",
                "Sibling",
                "OccurrenceOf",
            }
        )

        children = defaultdict(list)
        for child_edge in graph["edges"]["Child"]:
            f, t = child_edge[:2]
            children[f].append(t)

        def get_token_node(from_node_idx):
            to_visit = [from_node_idx]
            while len(to_visit) > 0:
                next_node = to_visit.pop()
                if next_node in token_sequence:
                    return next_node
                to_visit.extend(children[next_node])
            assert False, "No token found from the root node index."

        def backoff_to_parent(node_idx):
            return to_token_idx(next((k for k, v in children.items() if node_idx in v)))

        def to_token_idx(old_node_idx) -> int:
            if old_node_idx in old_node_id_to_new:
                return old_node_id_to_new[old_node_idx]
            to_visit = [old_node_idx]
            while len(to_visit) > 0:
                next_node_idx = to_visit.pop()
                next_node_label = graph["nodes"][next_node_idx]

                if next_node_label == "ComparisonTarget":
                    for child_node_idx in children[next_node_idx]:
                        child_node_label = graph["nodes"][child_node_idx]
                        if child_node_label in {
                            "<",
                            "<=",
                            "==",
                            "!=",
                            ">",
                            ">=",
                            "is",
                            "in",
                            "not",
                        }:
                            return old_node_id_to_new[child_node_idx]
                        elif child_node_label in {"IsNot", "NotIn"}:
                            # There are at least two tokens for these cases
                            return old_node_id_to_new[get_token_node(child_node_idx)]
                    else:
                        assert False, "One of the children, should be a comparison operator."
                elif next_node_label == "BinaryOperation":
                    for child_node_idx in children[next_node_idx]:
                        child_node_label = graph["nodes"][child_node_idx]
                        if child_node_label in {"+", "-", "*", "/", "//", "**", "%", "@", ">>", "<<", "|", "&", "^"}:
                            if child_node_idx not in old_node_id_to_new:
                                return backoff_to_parent(old_node_idx)
                            return old_node_id_to_new[child_node_idx]
                    else:
                        assert False, "One of the children, should be a binary operator."
                elif next_node_label in ("Assign", "AugAssign"):
                    for assign_target_node_idx in children[next_node_idx]:
                        if "=" in graph["nodes"][assign_target_node_idx]:
                            break
                    else:
                        assert False, "One child should always be an equals symbol."

                    return old_node_id_to_new[assign_target_node_idx]
                else:
                    for child_node_idx in children[next_node_idx]:
                        child_token_idx = old_node_id_to_new.get(child_node_idx)
                        if child_token_idx is not None:
                            return child_token_idx
                        else:
                            to_visit.append(child_node_idx)

            # Backoff to parent, should be rarely necessary (e.g. f-strings)
            return backoff_to_parent(old_node_idx)

        # Get ids of all AST nodes
        for child_edge in graph["edges"]["Child"]:
            f, t = child_edge[:2]
            old_node_id_to_new[f] = to_token_idx(f)
            old_node_id_to_new[t] = to_token_idx(t)

        # Map symbol nodes/edges to their first token occurrence
        orignal_symbol_node_idxs = defaultdict(list)
        for f, t in graph["edges"]["OccurrenceOf"]:
            orignal_symbol_node_idxs[t].append(f)
        for symbol_node_idx, target_tokens in orignal_symbol_node_idxs.items():
            old_node_id_to_new[symbol_node_idx] = min(old_node_id_to_new[t] for t in target_tokens)

        # Replace the edges
        new_edges = {}
        for edge_type, adj_list in graph["edges"].items():
            if edge_type in NON_TOKEN_EDGES:
                continue
            new_edges[edge_type] = [(to_token_idx(from_idx), to_token_idx(to_idx)) for from_idx, to_idx in adj_list]

        node_labels = graph["nodes"]
        all_tokens = [node_labels[i] for i in token_sequence]

        ref_token_idxs = []
        for ref_node_idx in graph["reference_nodes"]:
            ref_node_new_idx = to_token_idx(ref_node_idx)
            old_node_id_to_new[ref_node_idx] = ref_node_new_idx
            ref_token_idxs.append(ref_node_new_idx)

        return (
            all_tokens,
            old_node_id_to_new,
            new_edges,
            ref_token_idxs,
        )

    def __extract_token_sequence(self, graph: BugLabGraph) -> List[int]:
        next_token_edges = {f: t for f, t in graph["edges"]["NextToken"]}
        first_token = set(next_token_edges.keys()) - set(next_token_edges.values())
        if len(first_token) != 1:
            LOGGER.error(f"Encountered graph where the tokens are not connected in %s", graph["path"])
            return []
        first_token_idx = next(iter(first_token))

        def get_token_sequence():
            seen_tokens = {first_token_idx}
            next_token = first_token_idx
            yield next_token
            while next_token in next_token_edges:
                next_token = next_token_edges[next_token]
                if next_token in seen_tokens:
                    LOGGER.error("Cyclic token sequence in %s", graph["path"])
                    raise Exception()
                seen_tokens.add(next_token)
                yield next_token

        try:
            token_sequence = list(get_token_sequence())
        except:
            return []
        if len(token_sequence) != len(next_token_edges) + 1:
            LOGGER.error("Broken token sequence in %s", graph["path"])
            return []
        return token_sequence, next_token_edges

    def update_metadata_from(self, datapoint: BugLabData) -> None:
        try:
            token_data = self.__to_token_data(datapoint["graph"])
        except Exception as ex:
            LOGGER.exception("Error in generating token sequence for %s", datapoint["graph"]["path"], exc_info=ex)
            return None
        if token_data is None:
            return None
        node_labels, _, edges, _ = token_data

        for n in node_labels:
            self.__token_embedder.update_metadata_from(n)

        self.__edge_types.update(edges.keys())

    def finalize_metadata(self) -> None:
        self.__edge_types = list(self.__edge_types)
        self.__edge_type_to_idx = {t: i for i, t in enumerate(self.__edge_types)}

    def build_neural_module(self) -> SeqBugLabModule:
        return SeqBugLabModule(
            token_embedder=self.__token_embedder.build_neural_module(),
            embedding_dim=self.__token_embedder.embedding_size,
            num_edge_types=len(self.__edge_types),
            num_layers=self.__num_layers,
            num_heads=self.__num_heads,
            intermediate_dimension=self.__intermediate_dimension_size,
            dropout_rate=self.__dropout_rate,
            rewrite_vocabulary_size=len(self._target_rewrite_ops),
            layer_type=self.__layer_type,
            buggy_samples_weight_schedule=self.__buggy_samples_weight_schedule,
            generator_loss_type=self.__generator_loss_type,
            rezero_mode=self.__rezero_mode,
            normalisation_mode=self.__normalisation_mode,
        )

    def tensorize(self, datapoint: BugLabData) -> Optional[SeqModelTensorizedSample]:

        if "candidate_rewrite_logprobs" in datapoint:
            assert not self._tensorize_only_at_target_location_rewrites

        try:
            token_data = self.__to_token_data(datapoint["graph"])
        except Exception as ex:
            LOGGER.exception("Error in generating token sequence for %s", datapoint["graph"]["path"], exc_info=ex)
            return None
        if token_data is None:
            return None
        node_labels, old_node_id_to_new, edges, reference_nodes = token_data

        if len(node_labels) > self.__max_seq_size:
            LOGGER.debug("Rejecting sample with %s tokens.", len(node_labels))
            return None

        # Note: the transform_candidate_node_idxs defined below may contain duplicates. This is thanks
        # to the non-unique mapping of nodes from the graph to the sequence.
        candidate_node_idxs, inv = np.unique(datapoint["graph"]["reference_nodes"], return_inverse=True)
        transformed_candidate_node_idxs = np.array([old_node_id_to_new[n] for n in candidate_node_idxs])

        if datapoint["target_fix_action_idx"] is not None:
            target_node_idx = inv[datapoint["target_fix_action_idx"]]
            assert (
                old_node_id_to_new[datapoint["graph"]["reference_nodes"][datapoint["target_fix_action_idx"]]]
                == transformed_candidate_node_idxs[target_node_idx]
            )
        else:
            target_node_idx = None

        (
            # Rewrites
            target_rewrite_node_ids,
            target_rewrites,
            target_rewrite_to_location_group,
            correct_rewrite_target,
            text_rewrite_original_idx,
            # VarMisuse
            varmisused_node_ids,
            candidate_symbol_to_varmisused_location,
            candidate_symbol_node_ids,
            correct_candidate_symbol_node,
            varmisuse_rewrite_original_idx,
            # Arg Swap
            call_node_ids,
            candidate_swapped_node_ids,
            correct_swapped_pair,
            swapped_pair_to_call,
            swapped_rewrite_original_ids,
            # All rewrites
            repr_location_group_ids,
        ) = self._compute_rewrite_data(datapoint, candidate_node_idxs)

        target_rewrite_node_ids = [old_node_id_to_new[n] for n in target_rewrite_node_ids]
        varmisused_node_ids = [old_node_id_to_new[n] for n in varmisused_node_ids]
        candidate_symbol_node_ids = [old_node_id_to_new[n] for n in candidate_symbol_node_ids]
        call_node_ids = [old_node_id_to_new[n] for n in call_node_ids]
        candidate_swapped_node_ids = [
            (old_node_id_to_new[p1], old_node_id_to_new[p2]) for p1, p2 in candidate_swapped_node_ids
        ]

        return SeqModelTensorizedSample(
            target_subtokens_ids=[self.__token_embedder.tensorize(t) for t in node_labels],
            intra_token_edges=edges,
            candidate_location_idxs=transformed_candidate_node_idxs,
            target_location_idx=target_node_idx,
            node_mappings=old_node_id_to_new,
            ### Bug repair
            # Text rewrite
            target_rewrite_node_ids=target_rewrite_node_ids,
            target_rewrites=target_rewrites,
            target_rewrite_to_location_group=target_rewrite_to_location_group,
            correct_rewrite_target=correct_rewrite_target,
            text_rewrite_original_idx=text_rewrite_original_idx,
            # Var Misuse
            varmisused_node_ids=varmisused_node_ids,
            candidate_symbol_node_ids=candidate_symbol_node_ids,
            candidate_symbol_to_varmisused_node=candidate_symbol_to_varmisused_location,
            correct_candidate_symbol_node=correct_candidate_symbol_node,
            candidate_rewrite_original_idx=varmisuse_rewrite_original_idx,
            # Arg Swap
            call_node_ids=call_node_ids,
            candidate_swapped_node_ids=candidate_swapped_node_ids,
            swapped_pair_to_call=swapped_pair_to_call,
            correct_swapped_pair=correct_swapped_pair,
            pair_rewrite_original_idx=swapped_rewrite_original_ids,
            # All rewrites
            num_rewrite_locations_considered=len(repr_location_group_ids),
            # Rewrite logprobs
            rewrite_logprobs=datapoint.get("candidate_rewrite_logprobs", None),
        )

    def initialize_minibatch(self) -> Dict[str, Any]:
        return {
            "input_subtoken_ids": [],
            "edges": [],
            "edge_types": [],
            "candidate_location_idxs": [],
            "has_bug": [],
            "target_location_idxs": [],
            "node_mappings": [],
            # Repair
            "mb_num_repair_groups": 0,
            "mb_num_rewrite_candidates": 0,
            # Text rewrites
            "target_rewrite_node_ids": [],
            "target_rewrites": [],
            "rewrite_to_location_group": [],
            "correct_rewrite_idxs": [],
            "text_rewrite_original_idxs": [],
            "text_rewrite_idxs": [],
            # Var Misuse
            "varmisused_node_ids": [],
            "candidate_symbol_node_ids": [],
            "candidate_symbol_to_location_group": [],
            "correct_candidate_symbols": [],
            "candidate_rewrite_original_idxs": [],
            "candidate_rewrite_idxs": [],
            # Arg Swaps
            "call_node_ids": [],
            "candidate_swapped_node_ids": [],
            "swapped_pair_to_call_location_group": [],
            "correct_swapped_pair": [],
            "pair_rewrite_original_idx": [],
            "pair_rewrite_idxs": [],
            # Rewrite logprobs
            "rewrite_to_graph_id": [],
            "rewrite_logprobs": [],
            "no_bug_rewrite_logprobs": [],
        }

    def extend_minibatch_with(
        self, tensorized_datapoint: SeqModelTensorizedSample, partial_minibatch: Dict[str, Any]
    ) -> bool:
        current_sample_idx = len(partial_minibatch["input_subtoken_ids"])
        partial_minibatch["input_subtoken_ids"].append(tensorized_datapoint.target_subtokens_ids)

        for edge_type, adj_list in tensorized_datapoint.intra_token_edges.items():
            partial_minibatch["edge_types"].extend(self.__edge_type_to_idx[edge_type] for _ in adj_list)
            partial_minibatch["edges"].extend((current_sample_idx, f, t) for f, t in adj_list)

        partial_minibatch["has_bug"].append(tensorized_datapoint.target_location_idx is not None)

        partial_minibatch["node_mappings"].append(tensorized_datapoint.node_mappings)

        target_location_idx = tensorized_datapoint.target_location_idx or 0
        partial_minibatch["target_location_idxs"].append(
            target_location_idx + len(partial_minibatch["candidate_location_idxs"])
        )
        partial_minibatch["candidate_location_idxs"].extend(
            (current_sample_idx, tid) for tid in tensorized_datapoint.candidate_location_idxs
        )

        mb_rewrite_candidate_offset = partial_minibatch["mb_num_rewrite_candidates"]
        num_rewrite_candidates_in_datapoint = 0
        mb_num_repair_groups_offset = partial_minibatch["mb_num_repair_groups"]

        # Text rewrites
        if tensorized_datapoint.correct_rewrite_target is not None:
            partial_minibatch["correct_rewrite_idxs"].append(
                tensorized_datapoint.correct_rewrite_target + len(partial_minibatch["target_rewrites"])
            )
        partial_minibatch["target_rewrite_node_ids"].extend(
            (current_sample_idx, t) for t in tensorized_datapoint.target_rewrite_node_ids
        )
        partial_minibatch["target_rewrites"].extend(tensorized_datapoint.target_rewrites)
        partial_minibatch["rewrite_to_location_group"].extend(
            t + mb_num_repair_groups_offset for t in tensorized_datapoint.target_rewrite_to_location_group
        )
        partial_minibatch["text_rewrite_idxs"].extend(
            t + mb_rewrite_candidate_offset for t in tensorized_datapoint.text_rewrite_original_idx
        )
        num_rewrite_candidates_in_datapoint += len(tensorized_datapoint.text_rewrite_original_idx)

        # Var misuse
        if tensorized_datapoint.correct_candidate_symbol_node is not None:
            partial_minibatch["correct_candidate_symbols"].append(
                tensorized_datapoint.correct_candidate_symbol_node
                + len(partial_minibatch["candidate_symbol_to_location_group"])
            )
        partial_minibatch["varmisused_node_ids"].extend(
            (current_sample_idx, t) for t in tensorized_datapoint.varmisused_node_ids
        )
        partial_minibatch["candidate_symbol_node_ids"].extend(
            (current_sample_idx, t) for t in tensorized_datapoint.candidate_symbol_node_ids
        )
        partial_minibatch["candidate_symbol_to_location_group"].extend(
            t + mb_num_repair_groups_offset for t in tensorized_datapoint.candidate_symbol_to_varmisused_node
        )
        partial_minibatch["candidate_rewrite_idxs"].extend(
            t + mb_rewrite_candidate_offset for t in tensorized_datapoint.candidate_rewrite_original_idx
        )
        num_rewrite_candidates_in_datapoint += len(tensorized_datapoint.candidate_rewrite_original_idx)

        # Arg Swaps
        if tensorized_datapoint.correct_swapped_pair is not None:
            partial_minibatch["correct_swapped_pair"].append(
                tensorized_datapoint.correct_swapped_pair
                + len(partial_minibatch["swapped_pair_to_call_location_group"])
            )
        partial_minibatch["call_node_ids"].extend((current_sample_idx, t) for t in tensorized_datapoint.call_node_ids)
        partial_minibatch["candidate_swapped_node_ids"].extend(
            (current_sample_idx, t1, t2) for t1, t2 in tensorized_datapoint.candidate_swapped_node_ids
        )
        partial_minibatch["swapped_pair_to_call_location_group"].extend(
            t + mb_num_repair_groups_offset for t in tensorized_datapoint.swapped_pair_to_call
        )
        partial_minibatch["pair_rewrite_idxs"].extend(
            t + mb_rewrite_candidate_offset for t in tensorized_datapoint.pair_rewrite_original_idx
        )
        num_rewrite_candidates_in_datapoint += len(tensorized_datapoint.pair_rewrite_original_idx)
        partial_minibatch["mb_num_rewrite_candidates"] += num_rewrite_candidates_in_datapoint

        partial_minibatch["mb_num_repair_groups"] += tensorized_datapoint.num_rewrite_locations_considered

        # Visualization info
        partial_minibatch["text_rewrite_original_idxs"].append(tensorized_datapoint.text_rewrite_original_idx)
        partial_minibatch["candidate_rewrite_original_idxs"].append(tensorized_datapoint.candidate_rewrite_original_idx)
        partial_minibatch["pair_rewrite_original_idx"].append(tensorized_datapoint.pair_rewrite_original_idx)

        # Rewrite logprobs
        partial_minibatch["rewrite_to_graph_id"].extend([current_sample_idx] * num_rewrite_candidates_in_datapoint)
        logprobs = tensorized_datapoint.rewrite_logprobs
        if logprobs is not None:
            partial_minibatch["rewrite_logprobs"].extend(logprobs[:-1])
            partial_minibatch["no_bug_rewrite_logprobs"].append(logprobs[-1])

        return True

    def finalize_minibatch(
        self, accumulated_minibatch_data: Dict[str, Any], device: Union[str, torch.device]
    ) -> Dict[str, Any]:
        max_len = max(len(t) for t in accumulated_minibatch_data["input_subtoken_ids"])
        max_subtokens = self.__token_embedder.max_num_subtokens
        num_samples = len(accumulated_minibatch_data["input_subtoken_ids"])
        input_sequence_ids = np.zeros((num_samples, max_len, max_subtokens), dtype=np.int32)
        input_seq_num_subtokens = np.ones((num_samples, max_len), dtype=np.int32)
        token_seq_lengths = np.empty(num_samples, dtype=np.int32)

        for i, seq in enumerate(accumulated_minibatch_data["input_subtoken_ids"]):
            seq_len = len(seq)
            token_seq_lengths[i] = seq_len
            for j, subtokens in enumerate(seq):
                num_subtokens = min(len(subtokens), max_subtokens)
                input_sequence_ids[i, j, :num_subtokens] = subtokens[:num_subtokens]
                input_seq_num_subtokens[i, j] = num_subtokens

        minibatch = {
            "input_sequence_ids": torch.tensor(input_sequence_ids, dtype=torch.int64, device=device),
            "input_seq_num_subtokens": torch.tensor(input_seq_num_subtokens, dtype=torch.int64, device=device),
            "token_sequence_lengths": torch.tensor(token_seq_lengths, dtype=torch.int64, device=device),
            "edges": torch.tensor(accumulated_minibatch_data["edges"], dtype=torch.int64, device=device),
            "edge_types": torch.tensor(accumulated_minibatch_data["edge_types"], dtype=torch.int64, device=device),
            "has_bug": torch.tensor(accumulated_minibatch_data["has_bug"], dtype=torch.bool, device=device),
            "target_location_idxs": torch.tensor(
                accumulated_minibatch_data["target_location_idxs"], dtype=torch.int64, device=device
            ),
            "candidate_location_idxs": torch.tensor(
                accumulated_minibatch_data["candidate_location_idxs"], dtype=torch.int64, device=device
            ),
            "node_mappings": accumulated_minibatch_data["node_mappings"],
            # Text Rewrite
            "target_rewrite_node_ids": torch.tensor(
                accumulated_minibatch_data["target_rewrite_node_ids"], dtype=torch.int64, device=device
            ),
            "target_rewrites": torch.tensor(
                accumulated_minibatch_data["target_rewrites"], dtype=torch.int64, device=device
            ),
            "rewrite_to_location_group": torch.tensor(
                accumulated_minibatch_data["rewrite_to_location_group"], dtype=torch.int64, device=device
            ),
            "correct_rewrite_idxs": torch.tensor(
                accumulated_minibatch_data["correct_rewrite_idxs"], dtype=torch.int64, device=device
            ),
            # Var Misuse
            "varmisused_node_ids": torch.tensor(
                accumulated_minibatch_data["varmisused_node_ids"], dtype=torch.int64, device=device
            ),
            "candidate_symbol_node_ids": torch.tensor(
                accumulated_minibatch_data["candidate_symbol_node_ids"], dtype=torch.int64, device=device
            ),
            "candidate_symbol_to_location_group": torch.tensor(
                accumulated_minibatch_data["candidate_symbol_to_location_group"], dtype=torch.int64, device=device
            ),
            "correct_candidate_symbols": torch.tensor(
                accumulated_minibatch_data["correct_candidate_symbols"], dtype=torch.int64, device=device
            ),
            # Arg Swap
            "call_node_ids": torch.tensor(
                accumulated_minibatch_data["call_node_ids"], dtype=torch.int64, device=device
            ),
            "candidate_swapped_node_ids": torch.tensor(
                accumulated_minibatch_data["candidate_swapped_node_ids"], dtype=torch.int64, device=device
            ),
            "swapped_pair_to_call_location_group": torch.tensor(
                accumulated_minibatch_data["swapped_pair_to_call_location_group"], dtype=torch.int64, device=device
            ),
            "correct_swapped_pair": torch.tensor(
                accumulated_minibatch_data["correct_swapped_pair"], dtype=torch.int64, device=device
            ),
            # Visualisation info
            "text_rewrite_original_idxs": accumulated_minibatch_data["text_rewrite_original_idxs"],
            "candidate_rewrite_original_idxs": accumulated_minibatch_data["candidate_rewrite_original_idxs"],
            "pair_rewrite_original_idx": accumulated_minibatch_data["pair_rewrite_original_idx"],
            # Rewrite info
            "text_rewrite_idxs": torch.tensor(
                accumulated_minibatch_data["text_rewrite_idxs"], dtype=torch.int64, device=device
            ),
            "candidate_rewrite_idxs": torch.tensor(
                accumulated_minibatch_data["candidate_rewrite_idxs"], dtype=torch.int64, device=device
            ),
            "pair_rewrite_idxs": torch.tensor(
                accumulated_minibatch_data["pair_rewrite_idxs"], dtype=torch.int64, device=device
            ),
            # Rewrite logprobs
            "rewrite_to_graph_id": torch.tensor(
                accumulated_minibatch_data["rewrite_to_graph_id"], dtype=torch.int64, device=device
            ),
        }

        # Rewrite logprobs
        if accumulated_minibatch_data["rewrite_logprobs"]:
            rewrite_logprobs = (
                accumulated_minibatch_data["rewrite_logprobs"] + accumulated_minibatch_data["no_bug_rewrite_logprobs"]
            )
            minibatch["rewrite_logprobs"] = torch.tensor(rewrite_logprobs, dtype=torch.float32, device=device)
            minibatch["rewrite_to_graph_id"] = torch.tensor(
                accumulated_minibatch_data["rewrite_to_graph_id"], dtype=torch.int64, device=device
            )
        return minibatch

    def predict(
        self, data: Iterator[BugLabData], trained_nn: SeqBugLabModule, device, parallelize: bool
    ) -> Iterator[Tuple[BugLabData, Dict[int, float], List[float]]]:
        trained_nn.eval()
        with torch.no_grad(), self._tensorize_all_location_rewrites():
            for mb_data, original_datapoints in self.minibatch_iterator(
                self.tensorize_dataset(data, return_input_data=True, parallelize=parallelize),
                device,
                max_minibatch_size=50,
                parallelize=parallelize,
            ):
                num_samples = mb_data["input_sequence_ids"].shape[0]

                output_representation = trained_nn._compute_output_representation(
                    mb_data["input_sequence_ids"],
                    mb_data["input_seq_num_subtokens"],
                    mb_data["token_sequence_lengths"],
                    mb_data["edges"],
                    mb_data["edge_types"],
                )

                candidate_location_idxs = mb_data["candidate_location_idxs"]

                target_location_representations = output_representation[
                    candidate_location_idxs[:, 0], candidate_location_idxs[:, 1]
                ]  # [num_candidate_points, H]

                # Quick sanity check:
                # torch.exp(torch_scatter.scatter_logsumexp(candidate_log_probs, candidate_to_sample_idx))
                # should be ones.
                candidate_to_sample_idx, candidate_log_probs = trained_nn._compute_localization_logprobs(
                    target_location_representations, candidate_location_idxs[:, 0], num_samples
                )

                # Repair
                arg_swap_logprobs, text_repair_logprobs, varmisuse_logprobs = trained_nn._compute_repair_logprobs(
                    output_representation,
                    mb_data["target_rewrite_node_ids"],
                    mb_data["target_rewrites"],
                    mb_data["rewrite_to_location_group"],
                    mb_data["varmisused_node_ids"],
                    mb_data["candidate_symbol_node_ids"],
                    mb_data["candidate_symbol_to_location_group"],
                    mb_data["call_node_ids"],
                    mb_data["candidate_swapped_node_ids"],
                    mb_data["swapped_pair_to_call_location_group"],
                )

                candidate_to_sample_idx = candidate_to_sample_idx.cpu().numpy()
                candidate_log_probs = candidate_log_probs.cpu().numpy()

                yield from self._iter_per_sample_results(
                    mb_data,
                    candidate_to_sample_idx,
                    candidate_log_probs,
                    arg_swap_logprobs,
                    num_samples,
                    original_datapoints,
                    text_repair_logprobs,
                    varmisuse_logprobs,
                    node_mappings=mb_data["node_mappings"],
                )


import argparse
import os
import re
import time
from string import punctuation

import torch
import torch.nn as nn
import numpy as np
from pypinyin import pinyin, Style

import utils
import hparams as hp
import audio as Audio
from text import text_to_sequence
from model.fastspeech2 import FastSpeech2
from plot.utils import plot_mel

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

def read_source(source_path):
    ids = []
    sequences = []
    sentences = []
    lexicon = read_lexicon("text/pinyin-lexicon-r.txt")
    with open(source_path, "r") as f:
        for line in f:
            id_, sequence, sentence = line.strip("\n").split("|")
            ids.append(id_)

            sequences.append(np.array(text_to_sequence(sequence)))
            sentences.append(sentence)
    return ids, sequences

def read_lexicon(lex_path):
    lexicon = {}
    with open(lex_path) as f:
        for line in f:
            temp = re.split(r"\s+", line.strip("\n"))
            word = temp[0]
            phones = temp[1:]
            if word.lower() not in lexicon:
                lexicon[word.lower()] = phones
    return lexicon

def get_FastSpeech2(step):
    checkpoint_path = os.path.join(
        hp.checkpoint_path, "checkpoint_{}.pth.tar".format(step)
    )

    speaker_num = len(utils.get_speaker_to_id())
    model = nn.DataParallel(FastSpeech2(speaker_num))
    model.load_state_dict(torch.load(checkpoint_path)["model"])
    model.requires_grad = False
    model.eval()
    return model

def synthesize(
    model,
    vocoder,
    d_vec,
    x_vec,
    adain,
    speaker,
    gst,
    texts,
    file_ids,
    prefix="",
):
    src_len = torch.from_numpy(np.array([len(t) for t in texts])).to(device)
    texts = torch.from_numpy(utils.pad_1D(texts)).to(device)
    d_vec = (
        torch.from_numpy(np.array(d_vec))
        .to(device)
        .unsqueeze(0)
        .expand(len(file_ids), -1)
        if d_vec is not None
        else None
    )
    x_vec = (
        torch.from_numpy(np.array(x_vec))
        .to(device)
        .unsqueeze(0)
        .expand(len(file_ids), -1)
        if x_vec is not None
        else None
    )
    adain = (
        torch.from_numpy(np.array(adain))
        .to(device)
        .unsqueeze(0)
        .expand(len(file_ids), -1)
        if adain is not None
        else None
    )
    speakers = (
        torch.from_numpy(np.array([speaker])).to(device).expand(len(file_ids))
        if speaker is not None
        else None
    )
    gst = (
        torch.from_numpy(np.array(gst))
        .to(device)
        .unsqueeze(0)
        .expand(len(file_ids), -1)
        if gst is not None
        else None
    )

    (
        mel,
        mel_postnet,
        log_duration_output,
        duration_output,
        f0_output,
        energy_output,
        _,
        _,
        mel_len,
    ) = model(
        texts,
        src_len,
        max_src_len=torch.max(src_len).item(),
        d_vec=d_vec,
        x_vec=x_vec,
        adain=adain,
        speaker=speakers,
        use_gst=args.gst,
        gst=gst,
    )

    if not os.path.exists(hp.test_path):
        os.makedirs(hp.test_path)

    utils.vocoder_infer(
        mel_postnet.transpose(1, 2),
        vocoder,
        [
            os.path.join(hp.test_path, "{}_{}.wav".format(prefix, file_id))
            for file_id in file_ids
        ],
        mel_len * hp.hop_length,
    )

    for i in range(len(texts)):
        file_id = file_ids[i]
        src_length = src_len[i]
        mel_length = mel_len[i]
        mel_postnet_ = (
            mel_postnet[i, :mel_length].transpose(0, 1).detach().cpu().numpy()
        )
        f0_output_ = f0_output[i, :src_length].detach().cpu().numpy()
        energy_output_ = energy_output[i, :src_length].detach().cpu().numpy()
        duration_output_ = (
            duration_output[i, :src_length].detach().cpu().numpy().astype(np.int)
        )

        np.save(
            os.path.join(hp.test_path, "{}_{}.npy".format(prefix, file_id)),
            mel_postnet_.T,
        )

        plot_mel(
            [(mel_postnet_, f0_output_, energy_output_, duration_output_)],
            ["Synthesized Spectrogram"],
            filename=os.path.join(hp.test_path, "{}_{}.png".format(prefix, file_id)),
        )

if __name__ == "__main__":
    # Test
    parser = argparse.ArgumentParser()
    parser.add_argument("--step", type=int, default=500000)
    parser.add_argument(
        "--speaker",
        type=str,
        choices=[
            "TST_T1_S3",
            "TST_T1_S4",
            "TST_T1_S5",
            "TST_T2_S3",
            "TST_T2_S4",
            "TST_T2_S5",
        ],
    )
    parser.add_argument("--source", type=str)
    parser.add_argument("--d_vec", action="store_true")
    parser.add_argument("--x_vec", action="store_true")
    parser.add_argument("--adain", action="store_true")
    parser.add_argument("--speaker_emb", action="store_true")
    parser.add_argument("--gst", action="store_true")

    args = parser.parse_args()

    file_ids, texts = read_source(args.source)

    speaker_to_track = {
        "TST_T1_S3": "Track1_b",
        "TST_T1_S4": "Track1_b",
        "TST_T1_S5": "Track1_b",
        "TST_T2_S3": "Track2_b",
        "TST_T2_S4": "Track2_b",
        "TST_T2_S5": "Track2_b",
    }
    speaker_to_id = utils.get_speaker_to_id()

    prefix = "{}_{}".format(speaker_to_track[args.speaker], args.speaker[-2:])

    # Get averaged speaker embedding
    mel_path = os.path.join(hp.preprocessed_path, "mel")
    reference_mels = [
        np.load(os.path.join(mel_path, filename))
        for filename in os.listdir(mel_path)
        if args.speaker in filename
    ]

    d_vec = (
        np.mean(utils.get_pretrained_embedding(args.speaker, "d_vec"), 0)
        if args.d_vec
        else None
    )
    x_vec = (
        np.mean(utils.get_pretrained_embedding(args.speaker, "x_vec"), 0)
        if args.x_vec
        else None
    )
    adain = (
        np.mean(utils.get_pretrained_embedding(args.speaker, "adain_emb"), 0)
        if args.adain
        else None
    )

    model = get_FastSpeech2(args.step).to(device)
    vocoder = utils.get_vocoder()

    speaker = speaker_to_id[args.speaker] if args.speaker_emb else None
    gst = np.mean(utils.get_gst(reference_mels, model), 0) if args.gst else None

    with torch.no_grad():
        for i in range(0, len(file_ids), hp.batch_size):
            synthesize(
                model,
                vocoder,
                d_vec,
                x_vec,
                adain,
                speaker,
                gst,
                texts[i : min(len(texts), i + hp.batch_size)],
                file_ids[i : min(len(file_ids), i + hp.batch_size)],
                prefix,
            )


import os
import cv2
import numpy as np
from tinytag import TinyTag
from pytube import YouTube
import json
import sys
import progressbar

kernel = np.array([
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        ], np.uint8)

def create_line_drawing_image(img):
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_dilated = cv2.dilate(img_gray, kernel, iterations=1)
    img_diff = cv2.absdiff(img_dilated, img_gray)
    contour = 255 - img_diff
    return contour

def linetracevideo(inpt, output, directory):
    video = TinyTag.get(inpt)
    bitrate = str(video.bitrate)

    print('Separating audio...')
    os.system('ffmpeg -loglevel panic -i ' + inpt + ' ' + directory + 'mp3/audio.mp3')
    print('Done\n')

    print('Converting to linetraced images...')
    vidcap = cv2.VideoCapture(inpt)
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    fps = str(fps)
    success, image = vidcap.read()
    count = 1
    while success:
        img_contour = create_line_drawing_image(image)
        success, image = vidcap.read()
        cv2.imwrite(directory + "intermediate1/frame_%d.jpg" % count, img_contour)
        count += 1
    print('Done\n')

    print('Combining images...')
    os.system('ffmpeg -loglevel panic -r ' + fps + ' -i ' + directory +
              'intermediate1/frame_%01d.jpg -vcodec mpeg4 -y  -b:v ' + bitrate + 'k ' + directory +
              'intermediate2/movie.mp4 ')
    print('Done\n')

    print('Adding audio...')
    os.system('ffmpeg -loglevel panic -y -i ' + directory + 'mp3/audio.mp3 -r ' + fps + ' -i ' + directory +
              'intermediate2/movie.mp4 ' + output)
    print('Done\n')

    files = os.scandir(directory + 'intermediate1')
    for f in files:
        os.remove(f)
    files = os.scandir(directory + 'intermediate2')
    for f in files:
        os.remove(f)
    files = os.scandir(directory + 'mp3')
    for f in files:
        os.remove(f)
    main()

def linetraceimage(inpt, output):
    print('Linetracing image...')
    image = cv2.imread(inpt)
    img_contour = create_line_drawing_image(image)
    cv2.imwrite(output, img_contour)
    print('Done')
    main()

def linetrace():
    filelist = os.listdir('./input')
    midlist = ['back']
    for f in filelist:
        if ".mp4" in f or '.png' in f or '.jpg' in f or '.jpeg' in f:
            midlist.append(f)
    print("Press letter of item to linetrace OR type name of item")
    for i in range(len(midlist)):
        print(i, ":", midlist[i])

    choice = input()
    if int(choice) == 0:
        main()
    if len(choice) > 2:
        results = []
        print('\n' * 20)
        print('RESULTS FOR: ' + choice)
        for item in midlist:
            if choice.lower() in item.lower()[:-4]:
                results.append(item)
        print("Press letter of item to linetrace")
        for i in range(len(results)):
            print(i, ":", results[i])
        choice = input()
        file = ('./input/' + results[int(choice)])
        output = ('./output/' + results[int(choice)])
        directory = './'
        if file.lower().endswith(('.png', '.jpg', '.jpeg')):
            linetraceimage(file, output)
        else:
            linetracevideo(file, output, directory)
    else:
        file = ('./input/' + midlist[int(choice)])
        output = ('./output/' + midlist[int(choice)])
        directory = './'
        if file.lower().endswith(('.png', '.jpg', '.jpeg')):
            linetraceimage(file, output)
        else:
            linetracevideo(file, output, directory)

def progress(x, original):
    percent = (original - x) / original * 100
    bar = '['
    for i in range(int(percent)):
        bar += '='
    for i in range(int(100 - percent)):
        bar += '-'
    bar += ']'

    sys.stdout.write('\r')
    sys.stdout.write(bar + ' ' + str(int(percent)) + '%')
    sys.stdout.flush()

original = 0

def progress_func(a, b, c):
    global original
    if not progress_func.has_been_called:
        original = int(c)
    progress(int(c), original)
    progress_func.has_been_called = True

progress_func.has_been_called = False

def ytdl():
    link = input('Link: ')
    midlist = ['mp4', 'mp3']
    print("Press number of filetype to download")
    for i in range(len(midlist)):
        print(i, ":", midlist[i])
    extension = 'mp4'
    try:
        choice = int(input())
        if choice == 0:
            extension = 'mp4'
        elif choice == 1:
            extension = 'mp3'
        else:
            print("That isnt a choice!")
            ytdl()
    except:
        print('That isnt a number!')
        ytdl()
    try:
        yt = YouTube(link,
                     on_progress_callback=progress_func,
                     )
        if yt.age_restricted:
            yt = YouTube(
                link,
                use_oauth=True,
                allow_oauth_cache=True,
                on_progress_callback=progress_func,
            )
    except:
        raise Exception('Error downloading!')
    print('Title: ' + yt.title + '\nLength: ' + str(yt.length))
    download_folder = os.path.expanduser("~") + "/Downloads/"
    print('Downloading...')
    global ran
    ran = False
    yt.streams.filter(progressive=True, file_extension=extension).order_by('resolution').first().download(download_folder)
    print('\nSaved to: ' + download_folder)
    main()

def config():
    pass

def main():
    midlist = ['exit', 'config', 'ytdl', 'linetrace']
    print("Press number of tool to use")
    for i in range(len(midlist)):
        print(i, ":", midlist[i])
    kill = False
    try:
        choice = int(input())
        if choice == 0:
            kill = True
            exit(0)
        elif choice == 1:
            config()
        elif choice == 2:
            ytdl()
        elif choice == 3:
            linetrace()
        else:
            print('Not a choice!')
            main()
     except:
         if kill:
             exit(0)
         print('That isnt a number!')
         main()

if __name__ == '__main__':
    main()


#!/usr/bin/env python
import os
import numpy as np
import copy
from dataset import imagenet_vdet_classes
from proposal import vid_proposals
from ..utils.protocol import empty_det_from_box, score_proto, bbox_hash, det_score, boxes_at_frame, frame_path_at
from ..utils.common import imread
from ..utils.log import logging
from ..utils.timer import Timer
from ..utils.cython_nms import vid_nms
import heapq

def det_vid_with_box(vid_proto, box_proto, det_fun, net,
                    class_names=imagenet_vdet_classes):
    assert vid_proto['video'] == box_proto['video']
    root = vid_proto['root_path']
    det_proto = empty_det_from_box(box_proto)
    for frame in vid_proto['frames']:
        frame_id, path = frame['frame'], frame['path']
        det_cur_frame = \
            [i for i in det_proto['detections'] if i['frame'] == frame_id]
        if len(det_cur_frame) > 0:
            logging.info("Detecting in frame {}, {} boxes...".format(
                frame_id, len(det_cur_frame)))
            img = imread(os.path.join(root, path))
            boxes = [det['bbox'] for det in det_cur_frame]
            det_scores = det_fun(img, boxes, net)
            for det, scores in zip(det_cur_frame, det_scores):
                det['scores'] = score_proto(class_names, scores)
    return det_proto

def det_vid_without_box(vid_proto, det_fun, net,
                    class_names=imagenet_vdet_classes):
    logging.info("Generating proposals...")
    box_proto = vid_proposals(vid_proto)
    det_proto = det_vid_with_box(vid_proto, box_proto, det_fun, net,
                                 class_names)
    return det_proto

def det_vid_score(vid_proto, det_fun, net, box_proto=None,
                    class_names=imagenet_vdet_classes):
    if box_proto:
        return det_vid_with_box(vid_proto, box_proto, det_fun, net, class_names)
    else:
        return det_vid_without_box(vid_proto, det_fun, net, class_names)

def apply_vid_nms(det_proto, class_index, thres=0.3):
    logging.info('Apply NMS on video: {}'.format(det_proto['video']))
    new_det = {}
    new_det['video'] = det_proto['video']
    boxes = np.asarray([[det['frame'],]+det['bbox']+[det_score(det, class_index),]
             for det in det_proto['detections']], dtype='float32')
    keep = vid_nms(boxes, thresh=0.3)
    new_det['detections'] = copy.copy([det_proto['detections'][i] for i in keep])
    logging.info("{} / {} windows kept.".format(len(new_det['detections']),
                                         len(det_proto['detections'])))
    return new_det

def fast_rcnn_det_vid(net, vid_proto, box_proto, det_fun,
        class_names=imagenet_vdet_classes, max_per_image=100, thresh=0.05):
    """Test a Fast R-CNN network on a video protocol."""
    num_images = len(vid_proto['frames'])

    num_classes = len(class_names)
    # all detections are collected into:
    #    all_boxes[cls][image] = N x 5 array of detections in
    #    (x1, y1, x2, y2, score)
    all_boxes = [[[] for _ in xrange(num_images)]
                 for _ in xrange(num_classes)]

    # timers
    _t = {'im_detect' : Timer(), 'misc' : Timer()}

    for i, frame in enumerate(vid_proto['frames']):
        im = imread(frame_path_at(vid_proto, frame['frame']))
        _t['im_detect'].tic()
        orig_boxes = np.array([box['bbox'] for box in \
            boxes_at_frame(box_proto, frame['frame'])])
        scores, boxes = det_fun(net, im, orig_boxes)
        _t['im_detect'].toc()

        _t['misc'].tic()
        for j in xrange(1, num_classes):
            inds = np.where(scores[:, j] > thresh)[0]
            cls_scores = scores[inds, j]
            cls_boxes = boxes[inds, j*4:(j+1)*4]
            if len(cls_scores) > max_per_image:
                top_inds = np.argsort(-cls_scores)[:max_per_image]
                cls_scores = cls_scores[top_inds]
                cls_boxes = cls_boxes[top_inds, :]
            all_boxes[j][i] = \
                    np.hstack((cls_boxes, cls_scores[:, np.newaxis])) \
                    .astype(np.float32, copy=False)
        _t['misc'].toc()

        print 'im_detect: {:d}/{:d} {:.3f}s {:.3f}s' \
              .format(i + 1, num_images, _t['im_detect'].average_time,
                      _t['misc'].average_time)

    return all_boxes


import sys, os
from sklearn.metrics import roc_auc_score, recall_score, precision_score, roc_curve
import pandas as pd
import numpy as np
from evaluation.evaluateScoresList import evaluateScoresLists

DO_ROC_CURVE= False

# if INVERT_SELECTION==True, evaluate chains that are not in the file
INVERT_SELECTION= False
# if EVAL_AT_COMPLEX_LEVEL==True, select at complex level, otherwise select at chains levels
EVAL_AT_COMPLEX_LEVEL=True

pathDiffComplex= "/home/rsanchez/Tesis/rriPredMethod/data/bench5Data/benchmark5/difficultComplexes.txt"
difficultComplexes= set([])
with open(pathDiffComplex) as f:
  for line in f:
    prefix= line.split("_")[0]
    difficultComplexes.add(prefix)
pathMediumComplex= "/home/rsanchez/Tesis/rriPredMethod/data/bench5Data/benchmark5/mediumComplexes.txt"
with open(pathMediumComplex) as f:
  for line in f:
    prefix= line.split("_")[0]
    difficultComplexes.add(prefix)

def loadResults( resultsPath, fnameResults, selectChains, invertSelChains):
  prefix= fnameResults.split("_")[0].split(".")[0]
  if fnameResults.endswith(".lig"):
    chainType="l"
  else:
    chainType="r"

  scoresDf= pd.read_table(os.path.join(resultsPath, fnameResults), comment="#", sep="\s+", dtype={"resId":str, "chainId":str})
  if selectChains:
    if invertSelChains:
      scoresDf= scoresDf[~ scoresDf["chainId"].isin(selectChains) ]
    else:
      scoresDf= scoresDf[ scoresDf["chainId"].isin(selectChains) ]
  nPos= np.sum(scoresDf["categ"]==1 )
  if scoresDf.shape[0]==0:
    fraccPos= np.nan
  else:
    fraccPos= nPos/ float(scoresDf.shape[0])
  print(scoresDf.shape)
  return scoresDf, fraccPos

def get_single_chain_statistics(prefix, labels, scores):
  EVAL_PAIRS_AT= [ 2**3, 2**4, 2**5]
  precisionAt=[]
  recallAt=[]
  scores= np.array(scores)
  labels= np.array(labels)
  try:
    roc_complex= roc_auc_score(labels, scores)
  except ValueError:
    roc_complex= np.nan
  probability_sorted_indexes = scores.argsort(axis=0)
  probability_sorted_indexes = probability_sorted_indexes[::-1]
  for evalPoint in EVAL_PAIRS_AT:
    try:
      label_predictions= np.ones(scores.shape[0])* np.min( labels)
      label_predictions[probability_sorted_indexes[0 : evalPoint]]= np.repeat(1, evalPoint)
      precisionAt.append( precision_score(labels[probability_sorted_indexes[0 : evalPoint]],
                                        label_predictions[probability_sorted_indexes[0 : evalPoint]]))
      recallAt.append( recall_score(labels, label_predictions))
#      print(sum(labels==1), sum(label_predictions==1),precisionAt[-1], recallAt[-1])
    except IndexError:
      precisionAt.append( 0.0)
      recallAt.append( 0.0)
  summary= pd.DataFrame({"pdb":[prefix]})
  summary["auc_chains"]= [roc_complex]
  for evalPoint, precisionAt, recallAt in zip(EVAL_PAIRS_AT,precisionAt, recallAt):
    summary["prec_%d"%evalPoint]= [precisionAt]
    summary["reca_%d"%evalPoint]= [recallAt]

  rocCurve= None
  if DO_ROC_CURVE:
    fpr, tpr, __= roc_curve(labels, scores)
    rocCurve= (fpr, tpr, roc_complex)
  return summary, rocCurve

def loadChainsToEval(chainsToEvalFname):
  prefixToChains={}
  with open(chainsToEvalFname) as f:
    for line in f:
      lineArray= line.split()
      for elem in lineArray[1:]:
        prefix, chain= elem.split("_")
        if prefix not in prefixToChains:
          prefixToChains[prefix]=[]
        prefixToChains[prefix].append(chain)
  return prefixToChains

def getOptimThr(resultsPath, chainsToEvalFname):
  allScores=[]
  allLabels=[]
  perComplexSummaries=[]
  rocCurves= []
  prefixToChains= loadChainsToEval(chainsToEvalFname)
  prefixes= set([])
  nDifficult= 0
  fraccPos_list=[]
  for fname in sorted(os.listdir(resultsPath)):
    print(fname)
    if fname.endswith(".rec") or fname.endswith(".lig"):
      prefix = fname.split(".")[0]
      if not EVAL_AT_COMPLEX_LEVEL:
        selectChains= []
        if prefix in prefixToChains:
          selectChains= prefixToChains[prefix]
          if len(selectChains)==0:
            if not INVERT_SELECTION:
              continue
      else:
        selectChains= None
        if prefix in prefixToChains and INVERT_SELECTION:
          continue
        if not prefix in prefixToChains and not INVERT_SELECTION:
          continue

      results, fraccPos= loadResults( resultsPath, fname, selectChains, INVERT_SELECTION)
      if results is None or results.shape[0]==0: print("skip"); continue
      fraccPos_list.append(fraccPos)
      prefixes.add(prefix)
      if prefix in difficultComplexes:
        nDifficult+=1
        difficultComplexes.remove(prefix)

      scores= list(results["prediction"].values)
      labels= list(results["categ"].values)

      summary, rocCurve= get_single_chain_statistics(fname, labels, scores)
      if rocCurve: rocCurves.append(rocCurve)
      perComplexSummaries.append(summary)
#      print("%s %f"%(fname,roc_complex))
      allScores+= scores
      allLabels+= labels

  summary= pd.concat(perComplexSummaries, ignore_index=True)
  summary= summary.dropna(axis=0)
  means= summary.mean(axis=0)
  summary= summary.append( summary.ix[summary.shape[0]-1,:],ignore_index=True )
  summary.ix[summary.shape[0]-1,0]=  "mean"
  summary.ix[summary.shape[0]-1, 1:]=  means
  print("%d %d -> %f (difficult/total)"%( nDifficult, len(prefixes), float(nDifficult) / len(prefixes)) )
  print("fraccPos %f"%( np.nanmean(fraccPos_list) ))

  evaluateScoresLists(allLabels, allScores, summary, summary.iloc[-1,1], None if not DO_ROC_CURVE else rocCurves)

if __name__=="__main__":
  '''
python -m evaluation.evaluateChainsInList ~/Tesis/rriPredMethod/data/bench5Data/newCodeData/results_xgb/mixed_2/ ~/Tesis/rriPredMethod/data/partnerSpecificity/equivalentChainsB5.txt
  '''
  if len(sys.argv)!=3:
    raise ValueError("Error, bad number of arguments.usage:\npython -m evaluation.evaluateChainsInList pathToResults pathToChainsToEvaluate")
  resultsPath= os.path.expanduser(sys.argv[1])
  chainsPath= os.path.expanduser(sys.argv[2])
  getOptimThr(resultsPath, chainsPath)