import argparse, math, os, torch
import numpy as np
from tqdm import tqdm

from deeplab3.config.defaults import get_cfg_defaults
from dataloaders import make_data_loader
from deeplab3.modeling.sync_batchnorm.replicate import patch_replication_callback
from deeplab3.modeling import make_model
from deeplab3.utils.loss import SegmentationLosses
from deeplab3.utils.calculate_weights import calculate_weights_labels
from deeplab3.utils.lr_scheduler import LR_Scheduler
from deeplab3.utils.saver import Saver
from deeplab3.utils.summaries import TensorboardSummary
from deeplab3.evaluators import make_evaluator

class Trainer(object):
    def __init__(self, cfg):
        self.cfg = cfg

        # Define Saver
        self.saver = Saver(cfg)
        self.saver.save_experiment_config()
        # Define Tensorboard Summary
        self.summary = TensorboardSummary(self.saver.experiment_dir)
        self.writer = self.summary.create_summary()

        # Define Dataloader
        kwargs = {'num_workers': self.cfg.SYSTEM.NUM_WORKERS, 'pin_memory': True}
        self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(cfg, **kwargs)

        assert(self.nclass == cfg.DATASET.N_CLASSES)

        # Define network
        self.model = make_model(cfg)

        if self.cfg.TRAIN.FINETUNE:
            #Trains backbone less aggressively than final layers
            train_params = [{'params': self.model.get_1x_lr_params(), 'lr': self.cfg.TRAIN.LR},
                            {'params': self.model.get_10x_lr_params(), 'lr': self.cfg.TRAIN.LR * 10}]
        else:
            train_params = [{'params':self.model.parameters(), 'lr': self.cfg.TRAIN.LR}]

        # Define Optimizer
        self.optimizer = torch.optim.AdamW(train_params)

        # Define Criterion
        # whether to use class balanced weights
        if self.cfg.TRAIN.USE_BALANCED_WEIGHTS:
            classes_weights_path = os.path.join(cfg.DATASET.ROOT, cfg.DATASET.NAME+'_classes_weights.npy')
            if os.path.isfile(classes_weights_path):
                weight = np.load(classes_weights_path)
            else:
                weight = calculate_weights_labels(cfg.DATASET.ROOT, cfg.DATASET.NAME, self.train_loader, self.nclass)
            weight = torch.from_numpy(weight.astype(np.float32))
        else:
            weight = None
        self.criterion = SegmentationLosses(weight=weight, cuda=self.cfg.SYSTEM.CUDA).build_loss(mode=self.cfg.MODEL.LOSS_TYPE)

        # Define Evaluator
        self.evaluator = make_evaluator(self.cfg, self.nclass)

        # Using cuda
        if self.cfg.SYSTEM.CUDA:
            self.model = torch.nn.DataParallel(self.model, device_ids=self.cfg.SYSTEM.GPU_IDS)
            patch_replication_callback(self.model)
            self.model = self.model.cuda()

        # Resuming checkpoint
        self.best_pred = 0.0
        if self.cfg.CHECKPOINT.RESUME or self.cfg.TRAIN.FINETUNE:
            model_filepath = os.path.join(self.cfg.CHECKPOINT.DIRECTORY, self.cfg.CHECKPOINT.MODEL)

            if not os.path.isfile(model_filepath):
                raise RuntimeError("=> no checkpoint found at '{}'" .format(model_filepath))
            if cfg.SYSTEM.CUDA:
                checkpoint = torch.load(model_filepath, map_location=torch.device('cuda'))
            else:
                checkpoint = torch.load(model_filepath, map_location=torch.device('cpu'))
            self.cfg.TRAIN.START_EPOCH = checkpoint['epoch']

            strict = True
            for layer in cfg.CHECKPOINT.EXCLUDE:
                del checkpoint['state_dict'][layer]
                strict = False
            if self.cfg.SYSTEM.CUDA:
                self.model.module.load_state_dict(checkpoint['state_dict'], strict=strict)
            else:
                self.model.load_state_dict(checkpoint['state_dict'], strict=strict)

            # Load optimizer parameters and best previous prediction if resuming
            if self.cfg.CHECKPOINT.RESUME:
                self.optimizer.load_state_dict(checkpoint['optimizer'])
                self.best_pred = checkpoint['best_pred']

            # Clear start epoch if fine-tuning
            if self.cfg.TRAIN.FINETUNE:
                self.cfg.TRAIN.START_EPOCH = 0

            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(model_filepath, checkpoint['epoch']))

        # What is the accuracy before training?
        self.validation(self.cfg.TRAIN.START_EPOCH, self.cfg.TRAIN.START_EPOCH*len(self.train_loader))

    def training(self, epoch):
        train_loss = 0.0
        self.model.train()
        tbar = tqdm(self.train_loader)
        num_img_tr = len(self.train_loader)
        val_interval = math.floor(num_img_tr * self.cfg.TRAIN.EVAL_INTERVAL)
        for i, sample in enumerate(tbar):
            iter = i + num_img_tr * epoch
            if sample is not None:
                image, target = sample['image'], sample['label']
                if self.cfg.SYSTEM.CUDA:
                    image, target = image.cuda(), target.cuda()

                self.optimizer.zero_grad()
                try:
                    output = self.model(image)
                    loss = self.criterion(output, target)
                    loss.backward()
                    self.optimizer.step()
                    train_loss += loss.item()
                    tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
                    self.writer.add_scalar('train/total_loss_iter', loss.item(), iter)
                except ValueError as e:
                    print("{}: {}".format(str(e), sample['id']))

            if (i+1) % val_interval == 0:
                if not self.cfg.TRAIN.NO_VAL:
                    self.validation(epoch, iter)

        self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
        print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.cfg.TRAIN.BATCH_SIZE + image.data.shape[0]))
        print('Loss: %.3f' % train_loss)

        if self.cfg.TRAIN.NO_VAL:
            # save checkpoint every epoch
            is_best = False
            self.saver.save_checkpoint({
                'epoch': epoch + 1,
                'state_dict': self.model.module.state_dict(),
                'optimizer': self.optimizer.state_dict(),
                'best_pred': self.best_pred,
            }, is_best)

    def validation(self, epoch, iter):
        self.model.eval()
        self.evaluator.reset()

        max_val = min(self.cfg.TRAIN.VAL_MAX, len(self.val_loader))
        if max_val == -1:
            max_val = len(self.val_loader)
        tbar = tqdm(self.val_loader, desc='\r', total=max_val)
        test_loss = 0.0
        for i, sample in enumerate(tbar):
            image, target = sample['image'], sample['label']
            if self.cfg.SYSTEM.CUDA:
                image, target = image.cuda(), target.cuda()
            with torch.no_grad():
                output = self.model(image)
            loss = self.criterion(output, target)
            test_loss += loss.item()
            tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
            pred = output.data.cpu().numpy()
            target = target.cpu().numpy()
            pred = np.argmax(pred, axis=1)
            # Add batch sample into evaluator
            self.evaluator.add_batch(target, pred)

            if i==max_val:
                break

        # Fast test during the training
        new_pred = self.evaluator.write_metrics(self.writer, iter, i * self.cfg.TRAIN.BATCH_SIZE + image.data.shape[0])

        if new_pred > self.best_pred:
            is_best = True
            self.best_pred = new_pred
            self.saver.save_checkpoint({
                'epoch': epoch + 1,
                'state_dict': self.model.module.state_dict(),
                'optimizer': self.optimizer.state_dict(),
                'best_pred': self.best_pred,
            }, is_best)

def main():
    parser = argparse.ArgumentParser(description="PyTorch DeeplabV3Plus Training")
    parser.add_argument('config_file', help='config file path')
    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )

    args = parser.parse_args()

    cfg = get_cfg_defaults()
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    print(cfg)

    torch.manual_seed(cfg.SYSTEM.SEED)
    trainer = Trainer(cfg)

    for epoch in range(cfg.TRAIN.START_EPOCH, cfg.TRAIN.EPOCHS):
        trainer.training(epoch)

    trainer.writer.close()

if __name__ == "__main__":
   main()


import time
import os

import numpy
from absl import app, flags, logging
from absl.flags import FLAGS
import cv2
import numpy as np
import tensorflow as tf
from yolov3_tf2.models import (
    YoloV3, YoloV3Tiny
)
from yolov3_tf2.dataset import transform_images, load_tfrecord_dataset
from yolov3_tf2.utils import draw_outputs

flags.DEFINE_string('classes', './data/coco.names', 'path to classes file')
flags.DEFINE_string('weights', './checkpoints/yolov3.tf', 'path to weights file')
flags.DEFINE_boolean('tiny', False, 'yolov3 or yolov3-tiny')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_string('tfrecord', None, 'tfrecord instead of image')
flags.DEFINE_string('output', './2062_output.png', 'path to output image')
flags.DEFINE_integer('num_classes', 80, 'number of classes in the model')
flags.DEFINE_string('image', './data/images/test/2062.png', 'path to input image')

def main(_argv):
    # flags.DEFINE_string('image'+p, './data/images/'+p, 'path to input image')
    # flags.DEFINE_string('output', './output- '+p, 'path to output image')
    physical_devices = tf.config.experimental.list_physical_devices('GPU')
    for physical_device in physical_devices:
        tf.config.experimental.set_memory_growth(physical_device, True)

    if FLAGS.tiny:
        yolo = YoloV3Tiny(classes=FLAGS.num_classes)
    else:
        yolo = YoloV3(classes=FLAGS.num_classes)

    yolo.load_weights(FLAGS.weights).expect_partial()
    logging.info('weights loaded')

    class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
    logging.info('classes loaded')

    # if FLAGS.tfrecord:
    #     dataset = load_tfrecord_dataset(
    #         FLAGS.tfrecord, FLAGS.classes, FLAGS.size)
    #     dataset = dataset.shuffle(512)
    #     img_raw, _label = next(iter(dataset.take(1)))
    # else:
    path = 'data/images/validation'
    for p in os.listdir(path):
        img_full_path = os.path.join(path, p)
        print(p)
        img_raw = tf.image.decode_image(open(img_full_path, 'rb').read(), channels=3)

        img = tf.expand_dims(img_raw, 0)
        img = transform_images(img, FLAGS.size)

        t1 = time.time()
        boxes, scores, classes, nums = yolo(img)
        t2 = time.time()
        logging.info('time: {}'.format(t2 - t1))
        ########################################################################
        logging.info('detections:')
        detected_objects = []
        for i in range(nums[0]):
            logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
                                               np.array(scores[0][i]),
                                               np.array(boxes[0][i])))
            logging.info('\t{}, {}, {}'.format(p, class_names[int(classes[0][i])], np.array(scores[0][i])))
            detected_objects.append(class_names[int(classes[0][i])])

        c = numpy.array(detected_objects)
        unique, counts = numpy.unique(c, return_counts=True)

        if "car" in unique and counts[int(np.where(unique == "car")[0])] > 5 or "motorbike" in unique and counts[
            int(np.where(unique == "motorbike")[0])] > 2 or "bus" in unique and counts[
            int(np.where(unique == "bus")[0])] > 2 or "truck" in unique and counts[
            int(np.where(unique == "motorbike")[0])] > 3:
            logging.info('\t{}, {}, {}'.format(p, "traffic", 1.0))

        if "bird" in unique:
            logging.info('\t{}, {}, {}'.format(p, "flying", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "nature", 1.0))

        if "cat" or "dog" or "horse" or "sheep" or "cow" or "elephant" or "bear" or "zebra" or "giraffe" in unique:
            logging.info('\t{}, {}, {}'.format(p, "animal", 1.0))
        if "backpack" or "handbag" or "tie" in unique:
            logging.info('\t{}, {}, {}'.format(p, "outfit", 1.0))
        if "suitcase" in unique and "person":
            logging.info('\t{}, {}, {}'.format(p, "tourist", 1.0))
        if "suitcase" in unique and "person":
            logging.info('\t{}, {}, {}'.format(p, "tourist", 1.0))
        if "aeroplane" in unique:
            logging.info('\t{}, {}, {}'.format(p, "flying", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "travel", 1.0))
        if "skateboard" or "surfboard" or "snowboard" or "skis" in unique:
            logging.info('\t{}, {}, {}'.format(p, "snow", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "cool", 1.0))
        if "laptop" or "cell phone" or "remote" in unique:
            logging.info('\t{}, {}, {}'.format(p, "technology", 1.0))
        if "keyboard" or "cell phone" in unique:
            logging.info('\t{}, {}, {}'.format(p, "communication", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "connection", 1.0))
        if "toilet" or "vase" in unique:
            logging.info('\t{}, {}, {}'.format(p, "indoors", 1.0))
        if "chair" or "sofa" or "bed" or "diningtable" or "tvmonitor" or "microwave" or "oven" or \
                "toaster" or "hair drier" or "refrigerator" or "toothbrush" in unique:
            logging.info('\t{}, {}, {}'.format(p, "furniture", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "home", 1.0))
        if "cell phone" in unique:
            logging.info('\t{}, {}, {}'.format(p, "telephone", 1.0))
        if "boat" in unique:
            logging.info('\t{}, {}, {}'.format(p, "river", 1.0))
        if "bus" in unique or "traffic light" in unique or "car" in unique or "parking meter" in unique:
            logging.info('\t{}, {}, {}'.format(p, "street", 1.0))
        if "bus" in unique or "traffic light" in unique or "car" in unique or "truck" in unique or "train" in unique or "fire hydrant" in unique or "stop sign" in unique:
            logging.info('\t{}, {}, {}'.format(p, "road", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "city", 1.0))
            logging.info('\t{}, {}, {}'.format(p, "outdoors", 1.0))
        if "stop sign" in unique:
            logging.info('\t{}, {}, {}'.format(p, "sign", 1.0))
        if "clock" in unique:
            logging.info('\t{}, {}, {}'.format(p, "time", 1.0))
        if "bus" in unique:
            logging.info('\t{}, {}, {}'.format(p, "travel", 1.0))
        if "bottle" in unique:
            logging.info('\t{}, {}, {}'.format(p, "water", 1.0))

        for obj_cnt in zip(unique, counts):
            if obj_cnt[0] == "person":
                if obj_cnt[1] > 0:
                    logging.info('\t{}, {}, {}'.format(p, "human", 1.0))
                elif obj_cnt[1] > 2:
                    logging.info('\t{}, {}, {}'.format(p, "people", 1.0))
                elif 3 < obj_cnt[1] < 6:
                    logging.info('\t{}, {}, {}'.format(p, "meeting", 1.0))
                    logging.info('\t{}, {}, {}'.format(p, "group", 1.0))
                elif obj_cnt[1] >= 6:
                    logging.info('\t{}, {}, {}'.format(p, "crowd", 1.0))

            if obj_cnt[0] == "car" and obj_cnt[1] > 4:
                logging.info('\t{}, {}, {}'.format(p, "rally", 1.0))
            if obj_cnt[0] == "aeroplane" and obj_cnt[1] > 3:
                logging.info('\t{}, {}, {}'.format(p, "aircraft", 1.0))
                logging.info('\t{}, {}, {}'.format(p, "airplane", 1.0))

        img = cv2.cvtColor(img_raw.numpy(), cv2.COLOR_RGB2BGR)
        img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
        out_img = "data/output/out_" + p
        cv2.imwrite(out_img, img)
        logging.info('output saved to: {}'.format(out_img))

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass

cv2.waitKey(0)
cv2.destroyAllWindows()


"""Processes an MD2 file and returns mesh data in
an easy to read and process format.
"""

import os
import math
import struct
from collections import namedtuple, OrderedDict

import numpy

class MD2( object ):
    # The MD2 identifier
    # 'IDP2'
    id = 'IDP2'

    # The expected MD2 version
    version = 8

    header_layout = namedtuple(
        'MD2_Header',
        [
            'ident',
            'version',
            'skin_width',
            'skin_height',
            'frame_size',
            'num_skins',
            'num_vertices',
            'num_st',
            'num_tris',
            'num_glcmds',
            'num_frames',
            'offset_skins',
            'offset_st',
            'offset_tris',
            'offset_frames',
            'offset_glcmds',
            'offset_end'
            ]
        )

    frame_layout = namedtuple(
        'MD2_Frame',
        [
            'name',
            'vertices',
            'normals'
            ]
        )

    triangle_layout = namedtuple(
        'MD2_Triangles',
        [
            'vertex_indices',
            'tc_indices'
            ]
        )

    # The MD2 normal look-up table
    normal_lookup_table = numpy.array(
        [
            [-0.525731, 0.000000, 0.850651 ],
            [-0.442863, 0.238856, 0.864188 ],
            [-0.295242, 0.000000, 0.955423 ],
            [-0.309017, 0.500000, 0.809017 ],
            [-0.162460, 0.262866, 0.951056 ],
            [ 0.000000, 0.000000, 1.000000 ],
            [ 0.000000, 0.850651, 0.525731 ],
            [-0.147621, 0.716567, 0.681718 ],
            [ 0.147621, 0.716567, 0.681718 ],
            [ 0.000000, 0.525731, 0.850651 ],
            [ 0.309017, 0.500000, 0.809017 ],
            [ 0.525731, 0.000000, 0.850651 ],
            [ 0.295242, 0.000000, 0.955423 ],
            [ 0.442863, 0.238856, 0.864188 ],
            [ 0.162460, 0.262866, 0.951056 ],
            [-0.681718, 0.147621, 0.716567 ],
            [-0.809017, 0.309017, 0.500000 ],
            [-0.587785, 0.425325, 0.688191 ],
            [-0.850651, 0.525731, 0.000000 ],
            [-0.864188, 0.442863, 0.238856 ],
            [-0.716567, 0.681718, 0.147621 ],
            [-0.688191, 0.587785, 0.425325 ],
            [-0.500000, 0.809017, 0.309017 ],
            [-0.238856, 0.864188, 0.442863 ],
            [-0.425325, 0.688191, 0.587785 ],
            [-0.716567, 0.681718,-0.147621 ],
            [-0.500000, 0.809017,-0.309017 ],
            [-0.525731, 0.850651, 0.000000 ],
            [ 0.000000, 0.850651,-0.525731 ],
            [-0.238856, 0.864188,-0.442863 ],
            [ 0.000000, 0.955423,-0.295242 ],
            [-0.262866, 0.951056,-0.162460 ],
            [ 0.000000, 1.000000, 0.000000 ],
            [ 0.000000, 0.955423, 0.295242 ],
            [-0.262866, 0.951056, 0.162460 ],
            [ 0.238856, 0.864188, 0.442863 ],
            [ 0.262866, 0.951056, 0.162460 ],
            [ 0.500000, 0.809017, 0.309017 ],
            [ 0.238856, 0.864188,-0.442863 ],
            [ 0.262866, 0.951056,-0.162460 ],
            [ 0.500000, 0.809017,-0.309017 ],
            [ 0.850651, 0.525731, 0.000000 ],
            [ 0.716567, 0.681718, 0.147621 ],
            [ 0.716567, 0.681718,-0.147621 ],
            [ 0.525731, 0.850651, 0.000000 ],
            [ 0.425325, 0.688191, 0.587785 ],
            [ 0.864188, 0.442863, 0.238856 ],
            [ 0.688191, 0.587785, 0.425325 ],
            [ 0.809017, 0.309017, 0.500000 ],
            [ 0.681718, 0.147621, 0.716567 ],
            [ 0.587785, 0.425325, 0.688191 ],
            [ 0.955423, 0.295242, 0.000000 ],
            [ 1.000000, 0.000000, 0.000000 ],
            [ 0.951056, 0.162460, 0.262866 ],
            [ 0.850651,-0.525731, 0.000000 ],
            [ 0.955423,-0.295242, 0.000000 ],
            [ 0.864188,-0.442863, 0.238856 ],
            [ 0.951056,-0.162460, 0.262866 ],
            [ 0.809017,-0.309017, 0.500000 ],
            [ 0.681718,-0.147621, 0.716567 ],
            [ 0.850651, 0.000000, 0.525731 ],
            [ 0.864188, 0.442863,-0.238856 ],
            [ 0.809017, 0.309017,-0.500000 ],
            [ 0.951056, 0.162460,-0.262866 ],
            [ 0.525731, 0.000000,-0.850651 ],
            [ 0.681718, 0.147621,-0.716567 ],
            [ 0.681718,-0.147621,-0.716567 ],
            [ 0.850651, 0.000000,-0.525731 ],
            [ 0.809017,-0.309017,-0.500000 ],
            [ 0.864188,-0.442863,-0.238856 ],
            [ 0.951056,-0.162460,-0.262866 ],
            [ 0.147621, 0.716567,-0.681718 ],
            [ 0.309017, 0.500000,-0.809017 ],
            [ 0.425325, 0.688191,-0.587785 ],
            [ 0.442863, 0.238856,-0.864188 ],
            [ 0.587785, 0.425325,-0.688191 ],
            [ 0.688191, 0.587785,-0.425325 ],
            [-0.147621, 0.716567,-0.681718 ],
            [-0.309017, 0.500000,-0.809017 ],
            [ 0.000000, 0.525731,-0.850651 ],
            [-0.525731, 0.000000,-0.850651 ],
            [-0.442863, 0.238856,-0.864188 ],
            [-0.295242, 0.000000,-0.955423 ],
            [-0.162460, 0.262866,-0.951056 ],
            [ 0.000000, 0.000000,-1.000000 ],
            [ 0.295242, 0.000000,-0.955423 ],
            [ 0.162460, 0.262866,-0.951056 ],
            [-0.442863,-0.238856,-0.864188 ],
            [-0.309017,-0.500000,-0.809017 ],
            [-0.162460,-0.262866,-0.951056 ],
            [ 0.000000,-0.850651,-0.525731 ],
            [-0.147621,-0.716567,-0.681718 ],
            [ 0.147621,-0.716567,-0.681718 ],
            [ 0.000000,-0.525731,-0.850651 ],
            [ 0.309017,-0.500000,-0.809017 ],
            [ 0.442863,-0.238856,-0.864188 ],
            [ 0.162460,-0.262866,-0.951056 ],
            [ 0.238856,-0.864188,-0.442863 ],
            [ 0.500000,-0.809017,-0.309017 ],
            [ 0.425325,-0.688191,-0.587785 ],
            [ 0.716567,-0.681718,-0.147621 ],
            [ 0.688191,-0.587785,-0.425325 ],
            [ 0.587785,-0.425325,-0.688191 ],
            [ 0.000000,-0.955423,-0.295242 ],
            [ 0.000000,-1.000000, 0.000000 ],
            [ 0.262866,-0.951056,-0.162460 ],
            [ 0.000000,-0.850651, 0.525731 ],
            [ 0.000000,-0.955423, 0.295242 ],
            [ 0.238856,-0.864188, 0.442863 ],
            [ 0.262866,-0.951056, 0.162460 ],
            [ 0.500000,-0.809017, 0.309017 ],
            [ 0.716567,-0.681718, 0.147621 ],
            [ 0.525731,-0.850651, 0.000000 ],
            [-0.238856,-0.864188,-0.442863 ],
            [-0.500000,-0.809017,-0.309017 ],
            [-0.262866,-0.951056,-0.162460 ],
            [-0.850651,-0.525731, 0.000000 ],
            [-0.716567,-0.681718,-0.147621 ],
            [-0.716567,-0.681718, 0.147621 ],
            [-0.525731,-0.850651, 0.000000 ],
            [-0.500000,-0.809017, 0.309017 ],
            [-0.238856,-0.864188, 0.442863 ],
            [-0.262866,-0.951056, 0.162460 ],
            [-0.864188,-0.442863, 0.238856 ],
            [-0.809017,-0.309017, 0.500000 ],
            [-0.688191,-0.587785, 0.425325 ],
            [-0.681718,-0.147621, 0.716567 ],
            [-0.442863,-0.238856, 0.864188 ],
            [-0.587785,-0.425325, 0.688191 ],
            [-0.309017,-0.500000, 0.809017 ],
            [-0.147621,-0.716567, 0.681718 ],
            [-0.425325,-0.688191, 0.587785 ],
            [-0.162460,-0.262866, 0.951056 ],
            [ 0.442863,-0.238856, 0.864188 ],
            [ 0.162460,-0.262866, 0.951056 ],
            [ 0.309017,-0.500000, 0.809017 ],
            [ 0.147621,-0.716567, 0.681718 ],
            [ 0.000000,-0.525731, 0.850651 ],
            [ 0.425325,-0.688191, 0.587785 ],
            [ 0.587785,-0.425325, 0.688191 ],
            [ 0.688191,-0.587785, 0.425325 ],
            [-0.955423, 0.295242, 0.000000 ],
            [-0.951056, 0.162460, 0.262866 ],
            [-1.000000, 0.000000, 0.000000 ],
            [-0.850651, 0.000000, 0.525731 ],
            [-0.955423,-0.295242, 0.000000 ],
            [-0.951056,-0.162460, 0.262866 ],
            [-0.864188, 0.442863,-0.238856 ],
            [-0.951056, 0.162460,-0.262866 ],
            [-0.809017, 0.309017,-0.500000 ],
            [-0.864188,-0.442863,-0.238856 ],
            [-0.951056,-0.162460,-0.262866 ],
            [-0.809017,-0.309017,-0.500000 ],
            [-0.681718, 0.147621,-0.716567 ],
            [-0.681718,-0.147621,-0.716567 ],
            [-0.850651, 0.000000,-0.525731 ],
            [-0.688191, 0.587785,-0.425325 ],
            [-0.587785, 0.425325,-0.688191 ],
            [-0.425325, 0.688191,-0.587785 ],
            [-0.425325,-0.688191,-0.587785 ],
            [-0.587785,-0.425325,-0.688191 ],
            [-0.688191,-0.587785,-0.425325 ]
        ],
        dtype = numpy.float
        )

    # List of frame types used by Quake 2
    # http://tfc.duke.free.fr/old/models/md2.htm
    # beware: this page has a typo for the crouch_pain start frame
    animations = OrderedDict([
        ('stand',                (0, 39, 9.0)),
        ('run',                  (40, 45, 10.0)),
        ('attack',               (46, 53, 10.0)),
        ('pain_a',               (54, 57, 7.0)),
        ('pain_b',               (58, 61, 7.0)),
        ('pain_c',               (62, 65, 7.0)),
        ('jump',                 (66, 71, 7.0)),
        ('flip',                 (72, 83, 7.0)),
        ('salute',               (84, 94, 7.0)),
        ('fallback',             (95, 111, 10.0)),
        ('wave',                 (112, 122, 7.0)),
        ('point',                (123, 134, 6.0)),
        ('crouch_stand',         (135, 153, 10.0)),
        ('crouch_walk',          (154, 159, 7.0)),
        ('crouch_attack',        (160, 168, 10.0)),
        ('crouch_pain',          (169, 172, 7.0)),
        ('crouch_death',         (173, 177, 5.0)),
        ('death_fallback',       (178, 183, 7.0)),
        ('death_fallforward',    (184, 189, 7.0)),
        ('death_fallbackslow',   (190, 197, 7.0)),
        ('boom',                 (198, 198, 5.0)),
        ])

    def __init__( self ):
        super( MD2, self ).__init__()

        self.header = None
        self.skins = None
        self.triangles = None
        self.tcs = None
        self.frames = None

    def load( self, filename ):
        """
        Reads the MD2 data from the existing
        specified filename.

        @param filename: the filename of the md2 file
        to load.
        """
        with open( filename, 'rb' ) as f:
            self.load_from_buffer( f )

    def load_from_buffer( self, f ):
        """
        Reads the MD2 data from a stream object.

        Can be called instead of load() if data
        is not present in a file.

        @param f: the stream object, usually a file.
        """
        # read all the data from the file
        self.header = self.read_header( f )
        self.skins = self.read_skins( f, self.header )
        self.tcs = self.read_texture_coordinates( f, self.header )
        self.triangles = self.read_triangles( f, self.header )
        self.frames = self.read_frames( f, self.header )

    @staticmethod
    def _load_block( stream, format, count ):
        """
        Convenience method used to load blocks of
        data using the python 'struct' object format.

        Loads 'count' blocks from the file, each block
        will have the python struct format defined by 'format'.
        This is handy for loading large blocks without having
        to manually iterate over it.

        @param stream: the file object.
        @param format: the python 'struct' format of the block.
        @param count: the number of blocks to load.
        """
        def chunks( data, size ):
            """
            Return a generator that yields next 'size' bytes from data
            """
            offset = 0
            while offset < len( data ):
                yield data[ offset: offset + size ]
                offset += size

        struct_length = struct.calcsize( format )
        total_length = struct_length * count
        data = stream.read( total_length )

        if len( data ) < total_length:
            raise ValueError( "MD2: Failed to read '%d' bytes" % (total_length) )

        return [ struct.unpack(format, chunk) for chunk in chunks(data, struct_length) ]

    @staticmethod
    def read_header( f ):
        """
        Reads the MD2 header information from the MD2 file.

        @param f: the file object.
        @return Returns an header_layout named tuple.
        """
        # read the header
        # header is made up of 17 signed longs
        # this first is the ID which is also a 4 byte string
        header = MD2.header_layout._make(
            MD2._load_block( f, '< 4s16l', 1 )[ 0 ]
            )

        if header.ident != MD2.id:
            raise ValueError(
                "MD2 identifier is incorrect, expected '%i', found '%i'" % (
                    MD2.id,
                    header.ident
                    )
                )
        if header.version != MD2.version:
            raise ValueError(
                "MD2 version is incorrect, expected '%i', found '%i'" % (
                    MD2.version,
                    header.version
                    )
                )

        return header

    @staticmethod
    def read_skins( f, header ):
        """
        Reads the skin filenames out of the MD2 header.

        @param f: the file object.
        @param header: the loaded MD2 header.
        @return: Returns a python list of skin filenames.
        The list is a 1D list of size header.num_skins.
        """
        # seek to the skins offset
        f.seek( header.offset_skins, os.SEEK_SET )

        # skins are stored as a list of 64 signed byte strings
        # each string is a path relative to /baseq2
        skin_struct = struct.Struct( '< %s' % ('64s' * header.num_skins) )

        # read the skins and convert to list
        # strip any \x00 characters while we're at it
        # because python gets confused by them
        return [ skin.rstrip('\x00') for skin in skin_struct.unpack( f.read( skin_struct.size ) ) ]

    @staticmethod
    def read_texture_coordinates( f, header ):
        """
        Reads the texture coordinates from the MD2 file.

        @param f: the file object.
        @param header: the loaded MD2 header.
        @return: Returns a numpy array containing the texture
        coordinates. Values are converted from original
        absolute texels (0->width,0->height) to openGL
        coordinates (0.0->1.0).
        The array is an Nx2 dimension array where N
        is header.num_st.
        """
        # seek to the skins offset
        f.seek( header.offset_st, os.SEEK_SET )

        # st's are stored in a contiguous array of 2 short values
        # TCs do NOT map directly to vertices.
        # 1 vertex can have multiple TCs (one TC for each poly)
        # TCs are composed of 2 signed shorts
        tcs = numpy.array(
            MD2._load_block( f, '< 2h', header.num_st ),
            dtype = numpy.float
            )
        tcs.shape = (-1, 2)

        # convert from texel values to 0->1 float range
        tcs /= [ float(header.skin_width), float(header.skin_height) ]
        return tcs

    @staticmethod
    def read_triangles( f, header ):
        """
        Reads the triangle information from the MD2 file.

        Triangle information includes the vertex and
        texture coordinate indices.

        @param f: the file object.
        @param header: the loaded MD2 header.
        @return: Returns an MD2 named tuple.
        The vertex and texture coordinate indices are
        arrays of Nx3 dimensions for vertices and Nx2 for
        texture coordinates where N is header.num_tris.
        """
        # seek to the triangles offset
        f.seek( header.offset_tris, os.SEEK_SET )

        # triangles are stored as 3 unsigned shorts for the vertex indices
        # and 3 unsigned shorts for the texture coordinates indices
        triangles = numpy.array(
            MD2._load_block( f, '< 6H', header.num_tris ),
            dtype = numpy.uint16
            )
        triangles.shape = (-1, 6)

        # extract the vertex indices and tcs
        vertex_indices = triangles[ : , :3 ]
        tc_indices = triangles[ : , 3: ]

        # md2 triangles are clock-wise, we need to change
        # them to counter-clock-wise
        vertex_indices[ :,[1,2] ] = vertex_indices[ :,[2,1] ]
        tc_indices[ :,[1,2] ] = tc_indices[ :,[2,1] ]

        vertex_indices = vertex_indices.flatten()
        tc_indices = tc_indices.flatten()

        return MD2.triangle_layout(
            vertex_indices,
            tc_indices
            )

    @staticmethod
    def read_frames( f, header ):
        """
        Reads all frames from the MD2 file.

        This function simply calls read_frame in a loop.

        @param f: the file object.
        @param header: the loaded MD2 header.
        @return returns a python list of frame_layout
        named tuples. The list will be of length
        header.num_frames.
        """
        # seek to the frames offset
        f.seek( header.offset_frames, os.SEEK_SET )
        return [ MD2.read_frame( f, header ) for x in xrange( header.num_frames ) ]

    @staticmethod
    def read_frame( f, header ):
        """
        Reads a frame from the MD2 file.

        The stream must already be at the start of the
        frame.

        @param f: the file object.
        @param header: the loaded MD2 header.
        @return: Returns an frame_layout named tuple.
        Returned vertices and normals are as read from
        the file and are not ready to render.
        To render these must be ordered according to
        the indices specified in the triangle information.

        @see convert_indices_for_all_frames
        @see convert_indices_for_frame
        """
        # frame scale and translation are 2x3 32 bit floats
        frame_translations = numpy.array(
            MD2._load_block( f, '< 3f', 2 ),
            dtype = numpy.float
            )
        # extract the scale and translation vector
        scale = frame_translations[ 0 ]
        translation = frame_translations[ 1 ]

        # read the frame name
        # frame name is a 16 unsigned byte string
        name, = MD2._load_block( f, '< 16s', 1 )[0]
        # remove any \x00 characters as they confuse python
        name = name.strip( '\x00' )

        # frame has 3 unsigned bytes for the vertex coordinates
        # and 1 unsigned byte for the normal index
        frame_vertex_data = numpy.array(
            MD2._load_block( f, '<4B', header.num_vertices ),
            dtype = numpy.uint8
            )
        frame_vertex_data.shape = (-1, 4)

        # extract the vertex values
        vertices_short = frame_vertex_data[ :, :3 ]
        vertices = vertices_short.astype( numpy.float )
        vertices.shape = (-1, 3)

        # apply the frame translation
        vertices *= scale
        vertices += translation

        # re-orient the mesh
        # md2's have +Z as up, +Y as left, +X as forward
        # up: +Z, left: +Y, forward: +X
        # we want
        # up: +Y, left: -X, forward: -Z
        vertices[:,0],vertices[:,1],vertices[:,2] = \
            -vertices[:,1],vertices[:,2],-vertices[:,0]

        # extract the normal values
        normal_indices = frame_vertex_data[ :, 3 ]
        # convert from normal indice to normal vector
        normals = MD2.normal_lookup_table[ normal_indices ]
        normals.shape = (-1, 3)

        return MD2.frame_layout(
            name,
            vertices,
            normals
            )


#!/usr/bin/env python

# Psudo code:
'''
Create empty set of objects to track
Create current frame propoals
Create next frame proposals, for both next and current frame

for each image:
  Check if new first frame annos avaliable
  for each new first frame anno object:
    Read in anno
    Convert to json format, with bb and mask
    Calculate ReID embedding and add to JSON
    Add first frame anno as final solution
    Add first frame anno into next info for current frame
    Assign score to 1
    Calculate warped first frame anno, mask, bb and ReID
    Add to list of proposals for next frame, and next_info in objects to track
  Curr frame proposals = next frame proposals from last frame
  Read in new curr frame propals and add to curr frame proposals list
  for each not first frame anno object:
    Calculate weightings for WarpScores based on previous Total score for the assingned track from previous
    for each curr frame proposal:
      Calculate ReID score to each track next_info
      Calculate InVReID score to each track next_info
      Calculate WarpScore to each track next_info
      Calculate InvWarpScore to each track next_info
      Calculate Total Score for each proposal to each track
    Choose proposals for each track with argmax
    # maybe: choose blank proposal if score is too low
  Overlay proposals on one image choosing proposals with highest score
  From new overlaid proposals calculate bounding boxes
  Calculate new refined mask from new bounding boxes
  Overlay new refined masks again, again choosing that with the best score
  Add these to final output results
  Calculate warped mask into next frame, extract bb
  Calculate new refined warped mask, and ReID
  Add to list of proposals for next frame with current ReID
  Add to next info in objects to track with first frame ReID (explicit copy proabably needed)
'''

print("Starting up")
import glob
import numpy as np
import os
from copy import deepcopy as copy
from scipy.misc import imread

from MergeTrack.merge_functions import read_ann,read_props,calculate_scores,calculate_selected_props,\
  remove_mask_overlap,warp_proposals,update_templates,save_pngs,viz_scores,eval_video,calculate_gt_scores,\
  probabilitic_combination

from MergeTrack.refinement_net_functions import refinement_net_init, do_refinement
refinement_net = refinement_net_init()

from MergeTrack.ReID_net_functions import ReID_net_init, add_ReID
ReID_net = ReID_net_init()

input_images = "/content/test/JPEGImages/"
first_frame_anns = "/content/test/Annotations/"
input_proposals = "/content/content/PReMVOS/output/intermediate/ReID_proposals/"
input_optical_flow = "/content/drive/My Drive/flow/out/JPEGImages/"
output_images = "/content/drive/My Drive/flow/out/final/"

curr_run_num = 0
total_to_run = 1

def do_video(video_fn):

  print("Starting", video_fn)
  final_solution = []
  templates = []
  next_props = []
  image_fn_list = sorted(glob.glob(video_fn + "*"))
  print(image_fn_list)
  for image_id, image_fn in enumerate(image_fn_list):
    ann_fn = image_fn.replace(input_images, first_frame_anns).replace('.jpg', '.png')
    if glob.glob(ann_fn) and "00000.jpg" in image_fn:
      new_templates = read_ann(ann_fn)
      new_templates = add_ReID(new_templates, image_fn, ReID_net)
      templates = templates + copy(new_templates)
      next_props = next_props + copy(new_templates)
    if templates:
      prop_fn = image_fn.replace(input_images, input_proposals).replace('.jpg', '.json')
      proposals = next_props + read_props(prop_fn)

      all_scores = calculate_scores(proposals, templates)

      weighted_scores = np.dot(normalised_weights, all_scores.transpose((1, 0, 2)), )
      object_scores = np.dot(np.array([1, 1]), all_scores[:2, :, :].transpose((1, 0, 2)), )

      selected_props = calculate_selected_props(proposals, weighted_scores, templates, score_thesh, object_scores)
      selected_props = remove_mask_overlap(selected_props)

      optflow_fn = image_fn.replace(input_images, input_optical_flow).replace('.jpg', '.flo')
      try:
        if glob.glob(optflow_fn):
          next_image_fn = image_fn_list[image_id + 1]
          next_props = warp_proposals(selected_props, optflow_fn)

          next_props = do_refinement(next_props, next_image_fn, refinement_net)
          next_props = add_ReID(next_props, next_image_fn, ReID_net)
          templates = update_templates(templates, next_props)
        else:
          print("END")
      except:
        pass
      final_solution.append((image_fn, selected_props))
      output_image_fn = image_fn.replace(input_images, output_images).replace('.jpg', '.png')
      print(output_image_fn)
      save_pngs(selected_props, output_image_fn)
    else:
      output_image_fn = image_fn.replace(input_images, output_images).replace('.jpg', '.png')
      print(output_image_fn)
      img = imread(image_fn)
      temp_props = [dict()]
      temp_props[0]['mask'] = np.zeros_like(img[:, :, 1]).astype(np.uint8)
      save_pngs(temp_props, output_image_fn, empty=True)

score_thesh = 1e-10

weights = np.array([0.25920137, 0.22541801, 0.0775609,  0.12509281, 0.3127269 ])

normalised_weights = weights / np.sum(weights)

final_scores = []
video_list = sorted(glob.glob(input_proposals+"*/"))
video_list = [v.replace(input_proposals,input_images) for v in video_list]
num_per = int(np.ceil(len(video_list)/total_to_run))
print(curr_run_num,num_per,curr_run_num*num_per,(curr_run_num+1)*num_per)
for video_fn in video_list[curr_run_num*num_per : (curr_run_num+1)*num_per]:
  do_video(video_fn)


import sys
from pipeline import *
import argparse
import h5py
import os
import random
import time
import numpy as np
import torch
from torch.autograd import Variable
from torch import nn
from torch import cuda
from holder import *
from optimizer import *
from data import *
from util import *
from ema import *
from multiclass_loss import *

parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# data parameters
parser.add_argument('--dir', help="Path to the data dir", default="data/nli_aug/")
parser.add_argument('--train_data', help="Path to training data hdf5 file.", default="snli-train.hdf5")
parser.add_argument('--val_data', help="Path to validation data hdf5 file.", default="snli-val.hdf5")
parser.add_argument('--train_res', help="Path to training resource files, seperated by comma.", default="")
parser.add_argument('--val_res', help="Path to validation resource files, seperated by comma.", default="")
parser.add_argument('--word_vecs', help="The path to word embeddings", default = "glove.hdf5")
parser.add_argument('--char_idx', help="The path to word2char index file", default = "char.idx.hdf5")
parser.add_argument('--dict', help="The path to word dictionary", default = "snli.word.dict")
parser.add_argument('--char_dict', help="The path to char dictionary", default = "char.dict.txt")
parser.add_argument('--save_file', help="Path to where model to be saved.", default="model")
# training generic parameter
parser.add_argument('--percent', help="The percent of training data to use", type=float, default=1.0)
parser.add_argument('--epochs', help="The number of epoches for training", type=int, default=100)
parser.add_argument('--optim', help="The name of optimizer to use for training", default='adam')
parser.add_argument('--param_init_type', help="The type of parameter initialization", default='xavier_uniform')
parser.add_argument('--param_init', help="The scale of the normal distribution from which weights are initialized", type=float, default=0.01)
parser.add_argument('--clip_epoch', help="The starting epoch to enable clip", type=int, default=1)
parser.add_argument('--clip', help="The norm2 threshold to clip, set it to negative to disable", type=float, default=-1.0)
parser.add_argument('--ema', help="Whether to use EMA", type=int, default=0)
parser.add_argument('--mu', help="The mu ratio used in EMA", type=float, default=0.999)
parser.add_argument('--adam_betas', help="The betas used in adam", default='0.9,0.999')
parser.add_argument('--learning_rate', help="The learning rate for training", type=float, default=0.001)
parser.add_argument('--fix_word_vecs', help="Whether to make word embeddings NOT learnable", type=int, default=1)
parser.add_argument('--dropout', help="The dropout probability", type=float, default=0.2)
parser.add_argument('--print_every', help="Print stats after this many batches", type=int, default=1000)
parser.add_argument('--seed', help="The random seed", type=int, default=3435)
parser.add_argument('--gpuid', help="The GPU index, if -1 then use CPU", type=int, default=-1)
parser.add_argument('--use_char_enc', help="Whether to use char encoding", type=int, default=0)
parser.add_argument('--acc_batch_size', help="The accumulative batch size, -1 to disable", type=int, default=-1)
# dimensionality
parser.add_argument('--char_filters', help="The list of filters for char cnn", default='5')
parser.add_argument('--num_char', help="The number of distinct chars", type=int, default=61)
parser.add_argument('--char_emb_size', help="The input char embedding dim", type=int, default=20)
parser.add_argument('--char_enc_size', help="The input char encoding dim", type=int, default=100)
parser.add_argument('--hidden_size', help="The general hidden size of the pipeline", type=int, default=200)
parser.add_argument('--cls_hidden_size', help="The hidden size of the classifier", type=int, default=200)
parser.add_argument('--word_vec_size', help="The input word embedding dim", type=int, default=300)
parser.add_argument('--token_l', help="The maximal token length", type=int, default=16)
## pipeline specs
parser.add_argument('--encoder', help="The type of encoder", default="rnn")
parser.add_argument('--attention', help="The type of attention", default="local")
parser.add_argument('--classifier', help="The type of classifier", default="local")
parser.add_argument('--rnn_layer', help="The number of layers of rnn encoder", type=int, default=1)
parser.add_argument('--rnn_type', help="What type of rnn to use, default lstm", default='lstm')
parser.add_argument('--birnn', help="Whether to use bidirectional rnn", type=int, default=1)
parser.add_argument('--num_att_label', help='The number of attention labels', type=int, default=1)
parser.add_argument('--num_label', help="The number of prediction labels", type=int, default=3)
parser.add_argument('--constr_on', help="Directions of attentions to apply constraints on", default='1')
parser.add_argument('--cross_constr', help="The list of constraint layers to use, no if empty", default="")
parser.add_argument('--within_constr', help="The list of att constraint layers to use, no if empty", default="")
parser.add_argument('--rho_c', help="The weight of cross layer constraint", default='')
parser.add_argument('--rho_w', help="The weight of within layer constraint", type=float, default=1.0)

# train batch by batch, accumulate batches until the size reaches acc_batch_size
def train_epoch(opt, shared, m, optim, ema, data, epoch_id, sub_idx):
	train_loss = 0.0
	num_ex = 0
	start_time = time.time()
	num_correct = 0
	min_grad_norm2 = 1000000000000.0
	max_grad_norm2 = 0.0

	loss = MulticlassLoss(opt, shared)

	# subsamples of data
	# if subsample indices provided, permutate from subsamples
	#	else permutate from all the data
	data_size = sub_idx.size()[0]
	batch_order = torch.randperm(data_size)
	if sub_idx is not None:
		batch_order = sub_idx[batch_order]

	acc_batch_size = 0
	m.train(True)
	loss.begin_pass()
	m.begin_pass()
	for i in range(data_size):
		(data_name, source, target, char_source, char_target,
			batch_ex_idx, batch_l, source_l, target_l, label, res_map) = data[batch_order[i]]

		wv_idx1 = Variable(source, requires_grad=False)
		wv_idx2 = Variable(target, requires_grad=False)
		cv_idx1 = Variable(char_source, requires_grad=False)
		cv_idx2 = Variable(char_target, requires_grad=False)
		y_gold = Variable(label, requires_grad=False)

		# update network parameters
		shared.epoch = epoch_id
		m.update_context(batch_ex_idx, batch_l, source_l, target_l, res_map)

		# forward pass
		output = m.forward(wv_idx1, wv_idx2, cv_idx1, cv_idx2)

		# loss
		batch_loss = loss(output, y_gold)

		# stats
		train_loss += float(batch_loss.data)
		num_ex += batch_l
		time_taken = time.time() - start_time
		acc_batch_size += batch_l

		# accumulate grads
		batch_loss.backward()

		# accumulate current batch until the rolled up batch size exceeds threshold or meet certain boundary
		if i == data_size-1 or acc_batch_size >= opt.acc_batch_size or (i+1) % opt.print_every == 0:
			grad_norm2 = optim.step(m, acc_batch_size)
			if opt.ema == 1:
				ema.step(m)

			# clear up grad
			m.zero_grad()
			acc_batch_size = 0

			# stats
			grad_norm2_avg = grad_norm2
			min_grad_norm2 = min(min_grad_norm2, grad_norm2_avg)
			max_grad_norm2 = max(max_grad_norm2, grad_norm2_avg)
			time_taken = time.time() - start_time
			loss_stats = loss.print_cur_stats()

			if (i+1) % opt.print_every == 0:
				stats = '{0}, Batch {1:.1f}k '.format(epoch_id+1, float(i+1)/1000)
				stats += 'Grad {0:.1f}/{1:.1f} '.format(min_grad_norm2, max_grad_norm2)
				stats += 'Loss {0:.4f} '.format(train_loss / num_ex)
				stats += loss.print_cur_stats()
				stats += 'Time {0:.1f}'.format(time_taken)
				print(stats)

	perf, extra_perf = loss.get_epoch_metric()

	m.end_pass()
	loss.end_pass()

	return perf, extra_perf, train_loss / num_ex, num_ex

def train(opt, shared, m, optim, ema, train_data, val_data):
	best_val_perf = 0.0
	test_perf = 0.0
	train_perfs = []
	val_perfs = []
	extra_perfs = []

	train_idx, train_num_ex = train_data.subsample(opt.percent)
	print('{0} examples sampled for training'.format(train_num_ex))
	print('for the record, the first 10 training batches are: {0}'.format(train_idx[:10]))
	# sample the same proportion from the dev set as well
	#	but we don't want this to be too small
	minimal_dev_num = max(int(train_num_ex * 0.1), 1000)
	val_idx, val_num_ex = val_data.subsample(opt.percent, minimal_num=minimal_dev_num)
	print('{0} examples sampled for dev'.format(val_num_ex))
	print('for the record, the first 10 dev batches are: {0}'.format(val_idx[:10]))

	start = 0
	for i in range(start, opt.epochs):
		train_perf, extra_train_perf, loss, num_ex = train_epoch(opt, shared, m, optim, ema, train_data, i, train_idx)
		train_perfs.append(train_perf)
		extra_perf_str = ' '.join(['{:.4f}'.format(p) for p in extra_train_perf])
		print('Train {0:.4f} All {1}'.format(train_perf, extra_perf_str))

		# evaluate
		#	and save if it's the best model
		val_perf, extra_val_perf, val_loss, num_ex = validate(opt, shared, m, val_data, val_idx)
		val_perfs.append(val_perf)
		extra_perfs.append(extra_val_perf)
		extra_perf_str = ' '.join(['{:.4f}'.format(p) for p in extra_val_perf])
		print('Val {0:.4f} All {1}'.format(val_perf, extra_perf_str))

		perf_table_str = ''
		cnt = 0
		print('Epoch  | Train | Val ...')
		for train_perf, extra_perf in zip(train_perfs, extra_perfs):
			extra_perf_str = ' '.join(['{:.4f}'.format(p) for p in extra_perf])
			perf_table_str += '{0}\t{1:.4f}\t{2}\n'.format(cnt+1, train_perf, extra_perf_str)
			cnt += 1
		print(perf_table_str)

		if val_perf > best_val_perf:
			best_val_perf = val_perf
			print('saving model to {0}'.format(opt.save_file))
			param_dict = m.get_param_dict()
			save_param_dict(param_dict, '{0}.hdf5'.format(opt.save_file))
			save_opt(opt, '{0}.opt'.format(opt.save_file))
			# save ema
			if opt.ema == 1:
				ema_param_dict = ema.get_param_dict()
				save_param_dict(ema_param_dict, '{0}.ema.hdf5'.format(opt.save_file))

		else:
			print('skip saving model for perf <= {0:.4f}'.format(best_val_perf))

def validate(opt, shared, m, val_data, val_idx):
	m.train(False)

	val_loss = 0.0
	num_ex = 0

	loss = MulticlassLoss(opt, shared)

	data_size = val_idx.size()[0]
	print('validating on the {0} batches...'.format(val_idx.size()[0]))

	m.begin_pass()
	for i in range(data_size):
		(data_name, source, target, char_source, char_target,
			batch_ex_idx, batch_l, source_l, target_l, label, res_map) = val_data[val_idx[i]]

		wv_idx1 = Variable(source, requires_grad=False)
		wv_idx2 = Variable(target, requires_grad=False)
		cv_idx1 = Variable(char_source, requires_grad=False)
		cv_idx2 = Variable(char_target, requires_grad=False)
		y_gold = Variable(label, requires_grad=False)

		# update network parameters
		m.update_context(batch_ex_idx, batch_l, source_l, target_l, res_map)

		# forward pass
		pred = m.forward(wv_idx1, wv_idx2, cv_idx1, cv_idx2)

		# loss
		batch_loss = loss(pred, y_gold)

		# stats
		val_loss += float(batch_loss.data)
		num_ex += batch_l

	perf, extra_perf = loss.get_epoch_metric()
	m.end_pass()

	return (perf, extra_perf, val_loss / num_ex, num_ex)

def main(args):
	opt = parser.parse_args(args)
	shared = Holder()

	#
	opt.train_data = opt.dir + opt.train_data
	opt.val_data = opt.dir + opt.val_data
	opt.train_res = '' if opt.train_res == ''  else ','.join([opt.dir + path for path in opt.train_res.split(',')])
	opt.val_res = '' if opt.val_res == ''  else ','.join([opt.dir + path for path in opt.val_res.split(',')])
	opt.word_vecs = opt.dir + opt.word_vecs
	opt.char_idx = opt.dir + opt.char_idx
	opt.dict = opt.dir + opt.dict
	opt.char_dict = opt.dir + opt.char_dict

	torch.manual_seed(opt.seed)
	if opt.gpuid != -1:
		torch.cuda.set_device(opt.gpuid)
		torch.cuda.manual_seed_all(opt.seed)

	print(opt)

	# build model
	m = Pipeline(opt, shared)
	optim = Optimizer(opt, shared)
	ema = EMA(opt, shared)

	m.init_weight()
	model_parameters = filter(lambda p: p.requires_grad, m.parameters())
	num_params = sum([np.prod(p.size()) for p in model_parameters])
	print('total number of trainable parameters: {0}'.format(num_params))
	if opt.gpuid != -1:
		m = m.cuda()

	# loading data
	train_res_files = None if opt.train_res == '' else opt.train_res.split(',')
	train_data = Data(opt, opt.train_data, train_res_files)
	val_res_files = None if opt.val_res == '' else opt.val_res.split(',')
	val_data = Data(opt, opt.val_data, val_res_files)

	print('{0} batches in train set'.format(train_data.size()))

	train(opt, shared, m, optim, ema, train_data, val_data)

if __name__ == '__main__':
	sys.exit(main(sys.argv[1:]))


"""CAPRI module."""
import os
import shlex
import shutil
import subprocess
import tempfile
from functools import partial
from math import sqrt
from pathlib import Path

import numpy as np
from Bio import Align
from Bio.Align import substitution_matrices
from Bio.Seq import Seq
from fccpy import read_pdb
from fccpy.contacts import get_intermolecular_contacts
from pdbtools import pdb_segxchain

from haddock import log
from haddock.libs.libontology import PDBFile
from haddock.libs.libpdb import split_by_chain

RES_TO_BE_IGNORED = ["SHA"]

PROT_RES = [
    "ALA",
    "ARG",
    "ASN",
    "ASP",
    "CYS",
    "GLN",
    "GLU",
    "GLY",
    "HIS",
    "ILE",
    "LEU",
    "LYS",
    "MET",
    "PHE",
    "PRO",
    "SER",
    "THR",
    "TRP",
    "TYR",
    "VAL",
    ]

DNA_RES = ["DA", "DC", "DT", "DG"]
# Backbone
PROT_ATOMS = ["C", "N", "CA", "O"]
# Bases
DNA_ATOMS = [
    "C5",
    "N9",
    "N2",
    "C8",
    "O2",
    "N4",
    "N7",
    "C7",
    "N1",
    "N6",
    "C2",
    "O4",
    "C6",
    "N3",
    "C4",
    "O6",
    ]

class CAPRI:
    """CAPRI class."""

    def __init__(
            self,
            reference,
            model_list,
            receptor_chain,
            ligand_chain,
            aln_method,
            path,
            **params,
            ):
        self.reference = reference
        self.model_list = model_list
        self.irmsd_dic = {}
        self.lrmsd_dic = {}
        self.ilrmsd_dic = {}
        self.fnat_dic = {}
        self.dockq_dic = {}
        self.atoms = get_atoms(model_list + [reference])
        self.r_chain = receptor_chain
        self.l_chain = ligand_chain
        self.path = path

        # TODO: For scoring we might need to get one alignment per model
        reference = str(reference)
        model = model_list[0].rel_path
        align_func = get_align(aln_method, **params)
        self.numbering_dic = align_func(reference, model, path)
        if not self.numbering_dic:
            raise CAPRIError("Could not align reference and model")

        # Load the models in the class
        for struct in model_list:
            _ = self.add_chain_from_segid(struct.rel_path)

    def irmsd(self, cutoff=5.0):
        """Calculate the I-RMSD."""
        # Identify reference interface
        ref_interface_resdic = self.identify_interface(self.reference, cutoff)

        # Load interface coordinates
        ref_coord_dic, _ = self.load_coords(
            self.reference, ref_interface_resdic, match=False
            )

        for model in self.model_list:

            mod_coord_dic, _ = self.load_coords(
                model, ref_interface_resdic, match=True
                )

            # Here _coord_dic keys are matched
            #  and formatted as (chain, resnum, atom)
            #  we will use atoms that are present in both
            P = []
            Q = []
            for k in ref_coord_dic.keys() & mod_coord_dic.keys():
                ref_xyz = ref_coord_dic[k]
                mod_xyz = mod_coord_dic[k]

                Q.append(ref_xyz)
                P.append(mod_xyz)

            Q = np.asarray(Q)
            P = np.asarray(P)
            # write_coords("model.pdb", P)
            # write_coords("ref.pdb", Q)

            Q = Q - self.centroid(Q)
            P = P - self.centroid(P)
            U = self.kabsch(P, Q)
            P = np.dot(P, U)
            i_rmsd = self.calc_rmsd(P, Q)
            # write_coords("model_aln.pdb", P)
            # write_coords("ref_aln.pdb", Q)

            self.irmsd_dic[model] = i_rmsd

        return self.irmsd_dic

    def lrmsd(self):
        """Calculate the L-RMSD."""
        ref_coord_dic, _ = self.load_coords(self.reference)

        for model in self.model_list:

            mod_coord_dic, _ = self.load_coords(model, match=True)

            Q = []
            P = []
            # Note: this MUST be sorted since we will use the indexes to
            #  separate between receptor and ligand coordinates
            instersection = sorted(ref_coord_dic.keys() & mod_coord_dic.keys())

            chain_ranges = {}
            for i, segment in enumerate(instersection):
                chain, _, _ = segment
                if chain not in chain_ranges:
                    chain_ranges[chain] = []
                chain_ranges[chain].append(i)

            chain_ranges = make_range(chain_ranges)
            r_start, r_end = chain_ranges[self.r_chain]
            l_start, l_end = chain_ranges[self.l_chain]

            for k in instersection:
                ref_xyz = ref_coord_dic[k]
                mod_xyz = mod_coord_dic[k]

                Q.append(ref_xyz)
                P.append(mod_xyz)

            Q = np.asarray(Q)
            P = np.asarray(P)

            # write_coord_dic("ref.pdb", ref_coord_dic)
            # write_coord_dic("model.pdb", mod_coord_dic)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            # # move to the origin
            Q = Q - self.centroid(Q)
            P = P - self.centroid(P)

            # get receptor coordinates
            Q_r = Q[r_start: r_end - 1]
            P_r = P[r_start: r_end - 1]

            # Center receptors and get rotation matrix
            # Q_r = Q_r - self.centroid(Q_r)
            # P_r = P_r - self.centroid(P_r)

            U_r = self.kabsch(P_r, Q_r)

            # Center complexes at receptor centroids
            Q = Q - self.centroid(Q_r)
            P = P - self.centroid(P_r)

            # Apply rotation to complex
            #  - complex are now aligned by the receptor
            P = np.dot(P, U_r)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            # Identify the ligand coordinates
            Q_l = Q[l_start: l_end - 1]
            P_l = P[l_start: l_end - 1]

            # write_coords("ref_l.pdb", Q_l)
            # write_coords("model_l.pdb", P_l)

            # Calculate the RMSD of the ligands
            l_rmsd = self.calc_rmsd(P_l, Q_l)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            self.lrmsd_dic[model] = l_rmsd

        return self.lrmsd_dic

    def ilrmsd(self, cutoff=10.0):
        """Calculate the Interface Ligand RMSD."""
        # Identify interface
        ref_interface_resdic = self.identify_interface(self.reference, cutoff)

        # Load interface coordinates
        ref_coord_dic, _ = self.load_coords(self.reference, match=False)

        ref_int_coord_dic, _ = self.load_coords(
            self.reference, ref_interface_resdic, match=False
            )

        for model in self.model_list:

            mod_coord_dic, _ = self.load_coords(model, match=True)

            mod_int_coord_dic, _ = self.load_coords(
                model, ref_interface_resdic, match=True
                )

            # write_coord_dic("ref.pdb", ref_int_coord_dic)
            # write_coord_dic("model.pdb", mod_int_coord_dic)

            # find atoms present in both interfaces
            Q_int = []
            P_int = []
            common_keys = ref_int_coord_dic.keys() & mod_int_coord_dic.keys()
            for k in sorted(common_keys):
                ref_xyz = ref_int_coord_dic[k]
                mod_xyz = mod_int_coord_dic[k]

                Q_int.append(ref_xyz)
                P_int.append(mod_xyz)

            Q_int = np.asarray(Q_int)
            P_int = np.asarray(P_int)

            # write_coords("ref.pdb", Q_int)
            # write_coords("model.pdb", P_int)

            # find atoms present in both molecules
            Q = []
            P = []
            intersection = sorted(ref_coord_dic.keys() & mod_coord_dic.keys())
            chain_ranges = {}
            for i, segment in enumerate(intersection):
                chain, _, _ = segment
                if chain not in chain_ranges:
                    chain_ranges[chain] = []
                chain_ranges[chain].append(i)

            chain_ranges = make_range(chain_ranges)
            l_start, l_end = chain_ranges[self.l_chain]

            for k in sorted(ref_coord_dic.keys() & mod_coord_dic.keys()):
                ref_xyz = ref_coord_dic[k]
                mod_xyz = mod_coord_dic[k]

                Q.append(ref_xyz)
                P.append(mod_xyz)

            Q = np.asarray(Q)
            P = np.asarray(P)

            # write_coords("ref.pdb", Q)
            # write_coords("model.pdb", P)

            # put system at origin
            Q_int = Q_int - self.centroid(Q_int)
            P_int = P_int - self.centroid(P_int)

            # # put interfaces at the origin
            # Q_int = Q_int - self.centroid(Q_int)
            # P_int = P_int - self.centroid(P_int)

            # find the rotation that minimizes the interface rmsd
            U_int = self.kabsch(P_int, Q_int)
            P_int = np.dot(P_int, U_int)

            # write_coords("ref.pdb", Q_int)
            # write_coords("model.pdb", P_int)

            # # move the system to the centroid of the interfaces
            Q = Q - self.centroid(Q)
            P = P - self.centroid(P)

            # write_coords("ref_1.pdb", Q)
            # write_coords("model_1.pdb", P)

            Q = Q - self.centroid(Q_int)
            P = P - self.centroid(P_int)

            # write_coords("ref_2.pdb", Q)
            # write_coords("model_2.pdb", P)

            # apply this rotation to the model
            #  - complexes are now aligned by the interfaces
            P = np.dot(P, U_int)

            # write_coords("ref_i.pdb", Q)
            # write_coords("model_i.pdb", P)

            # Calculate the rmsd of the ligand
            Q_l = Q[l_start: l_end - 1]
            P_l = P[l_start: l_end - 1]

            # write_coords("ref_l.pdb", Q_l)
            # write_coords("model_l.pdb", P_l)

            # this will be the interface-ligand-rmsd
            i_l_rmsd = self.calc_rmsd(P_l, Q_l)
            self.ilrmsd_dic[model] = i_l_rmsd

        return self.ilrmsd_dic

    def fnat(self, cutoff=5.0):
        """Calculate the frequency of native contacts."""
        ref_contacts = self.load_contacts(self.reference, cutoff)
        for model in self.model_list:
            model_contacts = self.load_contacts(model, cutoff)
            intersection = ref_contacts & model_contacts
            fnat = len(intersection) / float(len(ref_contacts))
            self.fnat_dic[model] = fnat
        return self.fnat_dic

    def dockq(self):
        """Calculate the DockQ metric."""
        for model in self.model_list:
            irmsd = self.irmsd_dic[model]
            fnat = self.fnat_dic[model]
            lrmsd = self.lrmsd_dic[model]
            dockq = (
                float(fnat)
                + 1 / (1 + (irmsd / 1.5) * (irmsd / 1.5))
                + 1 / (1 + (lrmsd / 8.5) * (lrmsd / 8.5))
                ) / 3
            self.dockq_dic[model] = dockq

        return self.dockq_dic

    def output(
            self,
            clt_threshold,
            sortby_key,
            sort_ascending,
            ):
        """Output the CAPRI results to a .tsv file."""
        self._output_ss(sortby_key, sort_ascending)
        self._output_clt(clt_threshold, sortby_key, sort_ascending,)

    def _output_ss(self, sortby_key, sort_ascending):
        output_l = []
        for model in self.model_list:
            data = {}
            # keep always "model" the first key
            data["model"] = model
            # create the empty rank here so that it will appear
            #  as the second column
            data["caprieval_rank"] = None
            data["score"] = model.score
            if model in self.irmsd_dic:
                data["irmsd"] = self.irmsd_dic[model]
            if model in self.fnat_dic:
                data["fnat"] = self.fnat_dic[model]
            if model in self.lrmsd_dic:
                data["lrmsd"] = self.lrmsd_dic[model]
            if model in self.ilrmsd_dic:
                data["ilrmsd"] = self.ilrmsd_dic[model]
            if model in self.dockq_dic:
                data["dockq"] = self.dockq_dic[model]
            # add cluster data
            data["cluster-id"] = model.clt_id
            data["cluster-ranking"] = model.clt_rank
            data["model-cluster-ranking"] = model.clt_model_rank
            # list of dictionaries
            output_l.append(data)

        output_fname = Path(self.path, "capri_ss.tsv")
        self._dump_file(
            output_l,
            output_fname,
            sortby_key,
            sort_ascending,
            )

    def _output_clt(
            self,
            clt_threshold,
            sortby_key,
            sort_ascending,
            ):
        """Output cluster-based results."""
        has_cluster_info = any(m.clt_id for m in self.model_list)
        if not has_cluster_info:
            return

        # get the cluster data
        clt_data = dict(((m.clt_rank, m.clt_id), []) for m in self.model_list)

        # add models to each cluster
        for model in self.model_list:
            clt_data[(model.clt_rank, model.clt_id)].append(model)

        output_l = []
        for element in clt_data:
            data = {}
            number_of_models_in_cluster = len(clt_data[element])
            # TODO: Refactor these ugly try/excepts
            try:
                score_array = [v.score
                               for v in clt_data[element][: clt_threshold]]
                score_mean, score_stdev = self._calc_stats(
                    score_array, clt_threshold)
            except KeyError:
                score_mean = float("nan")
                score_stdev = float("nan")

            try:
                irmsd_array = [
                    self.irmsd_dic[v]
                    for v in clt_data[element]
                    [: clt_threshold]]
                irmsd_mean, irmsd_stdev = self._calc_stats(
                    irmsd_array, clt_threshold)
            except KeyError:
                irmsd_mean = float("nan")
                irmsd_stdev = float("nan")

            try:
                fnat_array = [self.fnat_dic[v]
                              for v in clt_data[element][:clt_threshold]]
                fnat_mean, fnat_stdev = self._calc_stats(
                    fnat_array, clt_threshold)
            except KeyError:
                fnat_mean = float("nan")
                fnat_stdev = float("nan")

            try:
                lrmsd_array = [self.lrmsd_dic[v]
                               for v in clt_data[element]
                               [: clt_threshold]]
                lrmsd_mean, lrmsd_stdev = self._calc_stats(
                    lrmsd_array, clt_threshold)
            except KeyError:
                lrmsd_mean = float("nan")
                lrmsd_stdev = float("nan")
            try:
                dockq_array = [self.dockq_dic[v]
                               for v in clt_data[element]
                               [: clt_threshold]]
                dockq_mean, dockq_stdev = self._calc_stats(
                    dockq_array, clt_threshold)
            except KeyError:
                dockq_mean = float("nan")
                dockq_stdev = float("nan")

            data["cluster_rank"] = element[0]
            data["cluster_id"] = element[1]
            data["n"] = number_of_models_in_cluster
            if number_of_models_in_cluster < clt_threshold:
                # under-evaluated, the mean was divided by a value
                #  larger than the total number of models in the cluster
                data["under_eval"] = "yes"
            else:
                data["under_eval"] = "-"

            data["score"] = score_mean
            data["score_std"] = score_stdev
            data["irmsd"] = irmsd_mean
            data["irmsd_std"] = irmsd_stdev
            data["fnat"] = fnat_mean
            data["fnat_std"] = fnat_stdev
            data["lrmsd"] = lrmsd_mean
            data["lrmsd_std"] = lrmsd_stdev
            data["dockqn"] = dockq_mean
            data["dockq_std"] = dockq_stdev

            output_l.append(data)

        output_fname = Path(self.path, "capri_clt.tsv")
        info_header = "#" * 40 + os.linesep
        info_header += "# `caprieval` cluster-based analysis" + os.linesep
        info_header += "#" + os.linesep
        info_header += f"# > sortby_key={sortby_key}" + os.linesep
        info_header += f"# > sort_ascending={sort_ascending}" + os.linesep
        info_header += f"# > clt_threshold={clt_threshold}" + os.linesep
        info_header += "#" + os.linesep
        info_header += (
            "# NOTE: if under_eval=yes, it means that there were less models in"
            " a cluster than" + os.linesep
            )
        info_header += (
            "#    clt_threshold, thus these values were under evaluated."
            + os.linesep
            )
        info_header += (
            "#   You might need to tweak the value of clt_threshold or change"
            " some parameters" + os.linesep
            )
        info_header += (
            "#    in `clustfcc` depending on your analysis." + os.linesep
            )
        info_header += "#" + os.linesep
        info_header += "#" * 40

        self._dump_file(
            output_l,
            output_fname,
            sortby_key,
            sort_ascending,
            info_header=info_header,
            )

    def _dump_file(
            self,
            container,
            output_fname,
            sortby_key,
            sort_ascending,
            info_header="",
            ):

        # rank
        ranked_output_l = self._rank(
            container, key='score', ascending=True
            )

        # sort
        sorted_keys = self._sort(
            ranked_output_l, key=sortby_key, ascending=sort_ascending
            )

        header = "\t".join(ranked_output_l[0].keys())

        if info_header:
            header = info_header + os.linesep + header

        with open(output_fname, "w") as out_fh:
            out_fh.write(header + os.linesep)
            for idx, _ in sorted_keys:
                row_l = []
                for value in ranked_output_l[idx].values():
                    if isinstance(value, Path):
                        row_l.append(str(value))
                    elif isinstance(value, PDBFile):
                        row_l.append(str(value.rel_path))
                    elif isinstance(value, int):
                        row_l.append(f"{value}")
                    elif isinstance(value, str):
                        row_l.append(f"{value}")
                    elif value is None:
                        row_l.append("-")
                    else:
                        row_l.append(f"{value:.3f}")
                out_fh.write("\t".join(row_l) + os.linesep)

    @staticmethod
    def _sort(container, key, ascending):
        # Sort the column
        key_values = [(i, k[key]) for i, k in enumerate(container)]
        key_values.sort(key=lambda x: x[1], reverse=not ascending)
        return key_values

    @staticmethod
    def _rank(container, key, ascending):
        rankkey_values = [(i, k[key]) for i, k in enumerate(container)]
        rankkey_values.sort(key=lambda x: x[1], reverse=not ascending)
        for i, k in enumerate(rankkey_values, start=1):
            idx, _ = k
            container[idx]["caprieval_rank"] = i
        return container

    @staticmethod
    def _calc_stats(data, n):
        """Calculate the mean and stdev."""
        mean = sum(data) / n
        var = sum((x - mean) ** 2 for x in data) / n
        stdev = sqrt(var)
        return mean, stdev

    @staticmethod
    def identify_interface(pdb_f, cutoff=5.0):
        """Identify the interface."""
        if isinstance(pdb_f, PDBFile):
            pdb_f = pdb_f.rel_path
        pdb = read_pdb(pdb_f)

        interface_resdic = {}
        for atom_i, atom_j in get_intermolecular_contacts(pdb, cutoff):

            if atom_i.chain not in interface_resdic:
                interface_resdic[atom_i.chain] = []
            if atom_j.chain not in interface_resdic:
                interface_resdic[atom_j.chain] = []

            if atom_i.resid not in interface_resdic[atom_i.chain]:
                interface_resdic[atom_i.chain].append(atom_i.resid)
            if atom_j.resid not in interface_resdic[atom_j.chain]:
                interface_resdic[atom_j.chain].append(atom_j.resid)

        return interface_resdic

    @staticmethod
    def load_contacts(pdb_f, cutoff=5.0):
        """Load residue-based contacts."""
        con_list = []
        if isinstance(pdb_f, PDBFile):
            pdb_f = pdb_f.rel_path
        structure = read_pdb(pdb_f)
        for atom_i, atom_j in get_intermolecular_contacts(structure, cutoff):
            con = (atom_i.chain, atom_i.resid, atom_j.chain, atom_j.resid)
            con_list.append(con)
        return set(con_list)

    @staticmethod
    def calc_rmsd(V, W):
        """Calculate the RMSD from two vectors."""
        diff = np.array(V) - np.array(W)
        N = len(V)
        return np.sqrt((diff * diff).sum() / N)

    @staticmethod
    def kabsch(P, Q):
        """Find the rotation matrix using Kabsch algorithm."""
        # Covariance matrix
        P = np.array(P)
        Q = np.array(Q)
        C = np.dot(np.transpose(P), Q)
        # use SVD
        V, S, W = np.linalg.svd(C)
        d = (np.linalg.det(V) * np.linalg.det(W)) < 0.0
        if d:
            S[-1] = -S[-1]
            V[:, -1] = -V[:, -1]
        # Create Rotation matrix U
        U = np.dot(V, W)
        return U

    @staticmethod
    def centroid(X):
        """Get the centroid."""
        X = np.array(X)
        return X.mean(axis=0)

    @staticmethod
    def add_chain_from_segid(pdb_path):
        """Replace the chainID with the segID."""
        temp_f = tempfile.NamedTemporaryFile(delete=False, mode="w+t")
        with open(pdb_path) as fh:
            for line in list(pdb_segxchain.run(fh)):
                temp_f.writelines(line)
        temp_f.close()
        # REPLACE!
        new_pdb_path = shutil.move(temp_f.name, pdb_path)
        return new_pdb_path

    def load_coords(self, pdb_f, filter_resdic=None, match=False):
        """Load coordinates from PDB."""
        coord_dic = {}
        chain_dic = {}
        idx = 0
        if isinstance(pdb_f, PDBFile):
            pdb_f = pdb_f.rel_path
        with open(pdb_f, "r") as fh:
            for line in fh.readlines():
                if line.startswith("ATOM"):

                    atom_name = line[12:16].strip()
                    resname = line[17:20].strip()
                    chain = line[21]
                    resnum = int(line[22:26])

                    x = float(line[30:38])
                    y = float(line[38:46])
                    z = float(line[46:54])
                    coords = np.asarray([x, y, z])

                    if match:
                        try:
                            resnum = self.numbering_dic[chain][resnum]
                        except KeyError:
                            # this residue is not matched, and so it should
                            #  not be considered
                            # self.log(
                            #     f"WARNING: {chain}.{resnum}.{atom_name}"
                            #     " was not matched!"
                            #     )
                            continue

                    # identifier = f"{chain}.{resnum}.{atom_name}"
                    identifier = (chain, resnum, atom_name)

                    if atom_name not in self.atoms[resname]:
                        continue

                    if chain not in chain_dic:
                        chain_dic[chain] = []

                    if filter_resdic:
                        # Only retrieve coordinates from the filter_resdic
                        if (
                                chain in filter_resdic
                                and resnum in filter_resdic[chain]
                                ):
                            coord_dic[identifier] = coords
                            chain_dic[chain].append(idx)
                            idx += 1

                    else:
                        # retrieve everything
                        coord_dic[identifier] = coords
                        chain_dic[chain].append(idx)
                        idx += 1

        chain_ranges = {}
        for chain in chain_dic:
            min_idx = min(chain_dic[chain])
            max_idx = max(chain_dic[chain])
            chain_ranges[chain] = (min_idx, max_idx)

        return coord_dic, chain_ranges

class CAPRIError(Exception):
    """Raised when something goes wrong with the CAPRI class."""

    def __init__(self, msg=""):
        self.msg = msg
        super().__init__(self.msg)

def get_atoms(pdb_list):
    """Identify what is the molecule type of each PDB."""
    atom_dic = {}
    atom_dic.update(dict((r, PROT_ATOMS) for r in PROT_RES))
    atom_dic.update(dict((r, DNA_ATOMS) for r in DNA_RES))
    for pdb in pdb_list:
        if isinstance(pdb, PDBFile):
            pdb = pdb.rel_path
        with open(pdb) as fh:
            for line in fh.readlines():
                if line.startswith(("ATOM", "HETATM")):
                    resname = line[17:20].strip()
                    atom_name = line[12:16].strip()
                    element = line[76:78].strip()
                    if (
                            resname not in PROT_RES
                            and resname not in DNA_RES
                            and resname not in RES_TO_BE_IGNORED
                            ):
                        # its neither DNA nor protein, use the heavy atoms
                        # WARNING: Atoms that belong to unknown residues must
                        #  be bound to a residue name;
                        #   For example: residue NEP, also contains
                        #  CB and CG atoms, if we do not bind it to the
                        #  residue name, the next functions will include
                        #  CG and CG atoms in the calculations for all
                        #  other residue names
                        if element != "H":
                            if resname not in atom_dic:
                                atom_dic[resname] = []
                            if atom_name not in atom_dic[resname]:
                                atom_dic[resname].append(atom_name)
    return atom_dic

def pdb2fastadic(pdb_f):
    """Write the sequence as a fasta."""
    res_codes = dict(
        [
            ("CYS", "C"),
            ("ASP", "D"),
            ("SER", "S"),
            ("GLN", "Q"),
            ("LYS", "K"),
            ("ILE", "I"),
            ("PRO", "P"),
            ("THR", "T"),
            ("PHE", "F"),
            ("ASN", "N"),
            ("GLY", "G"),
            ("HIS", "H"),
            ("LEU", "L"),
            ("ARG", "R"),
            ("TRP", "W"),
            ("ALA", "A"),
            ("VAL", "V"),
            ("GLU", "E"),
            ("TYR", "Y"),
            ("MET", "M"),
            ("DA", "A"),
            ("DG", "G"),
            ("DC", "C"),
            ("DT", "T"),
            ]
        )
    seq_dic = {}
    with open(pdb_f) as fh:
        for line in fh.readlines():
            if line.startswith("ATOM"):
                res_num = int(line[22:26])
                res_name = line[17:20].strip()
                chain = line[21]
                if res_name in RES_TO_BE_IGNORED:
                    continue
                try:
                    one_letter = res_codes[res_name]
                except KeyError:
                    one_letter = "X"
                if chain not in seq_dic:
                    seq_dic[chain] = {}
                seq_dic[chain][res_num] = one_letter
    return seq_dic

def get_align(method, **kwargs):
    """Get the alignment function."""
    log.info(f"Using {method} alignment")
    if method == "structure":
        return partial(align_strct, lovoalign_exec=kwargs["lovoalign_exec"])
    elif method == "sequence":
        return partial(align_seq)
    else:
        available_alns = ("sequence", "structure")
        raise ValueError(
            f"Alignment method {method!r} not recognized. "
            f"Available options are {', '.join(available_alns)}"
            )

def align_strct(reference, model, output_path, lovoalign_exec=None):
    """Structuraly align and get numbering relationship."""
    if lovoalign_exec is None:
        log.error(
            "Structural alignment needs LovoAlign "
            "get it at github.com/m3g/lovoalign"
            )
        raise CAPRIError("Path to LovoAlign executable required.")

    if not lovoalign_exec:
        raise CAPRIError("lovoalign_exec parameter not defined ")

    if not os.access(lovoalign_exec, os.X_OK):
        raise CAPRIError(f"{lovoalign_exec!r} for LovoAlign is not executable")

    numbering_dic = {}
    protein_a_dic = dict(
        (str(e.stem).split("_")[-1], e) for e in split_by_chain(reference)
        )
    protein_b_dic = dict(
        (str(e.stem).split("_")[-1], e) for e in split_by_chain(model)
        )

    # check if chain ids match
    if protein_a_dic.keys() != protein_b_dic.keys():
        # TODO: Make this a clearer raise
        return numbering_dic

    for chain in protein_a_dic.keys():
        pa_seqdic = pdb2fastadic(protein_a_dic[chain])
        pb_seqdic = pdb2fastadic(protein_b_dic[chain])
        # logging.debug(f"Structurally aligning chain {chain}")
        numbering_dic[chain] = {}
        cmd = (
            f"{lovoalign_exec} -p1 {protein_a_dic[chain]} "
            f"-p2 {protein_b_dic[chain]} "
            f"-c1 {chain} -c2 {chain}"
            )

        # logging.debug(f"Command is: {cmd}")
        p = subprocess.run(shlex.split(cmd), capture_output=True, text=True)
        lovoalign_out = p.stdout.split(os.linesep)

        # we don"t need the splitted proteins anymore
        protein_a_dic[chain].unlink()
        protein_b_dic[chain].unlink()

        # find out where the alignment starts and ends
        alignment_pass = True
        for i, line in enumerate(lovoalign_out):
            if "SEQUENCE ALIGNMENT" in line:
                # there are 2 extra white lines after this header
                alignment_start_index = i + 2
            elif "FINAL" in line:
                # there are 2 extra white lines after this header
                alignment_end_index = i - 2
            elif "ERROR" in line:
                failed_pdb = line.split()[-1]
                _msg = (
                    f"LovoAlign could not read {failed_pdb} " "is it a ligand?"
                    )
                log.warning(_msg)
                alignment_pass = False

                for elem in [k for k in pa_seqdic[chain]]:
                    numbering_dic[chain][elem] = elem

        if not alignment_pass:
            # This alignment failed, move on to the next
            log.warning(
                f"Skipping alignment of chain {chain}, "
                "used sequential matching"
                )
            continue

        aln_l = lovoalign_out[alignment_start_index:alignment_end_index]

        # dump this alignment to a file
        aln_fname = Path(output_path, f"lovoalign_{chain}.aln")
        log.debug(f"Writing alignment to {aln_fname.name}")
        with open(aln_fname, "w") as fh:
            fh.write(os.linesep.join(aln_l))

        # remove the line between the alignment segments
        alignment = [aln_l[i: i + 3][:2] for i in range(0, len(aln_l), 3)]
        # 100% (5 identical nucleotides / min(length(A),length(B))).
        len_seq_a = len(pa_seqdic[chain])
        len_seq_b = len(pb_seqdic[chain])
        identity = (
            (len_seq_a - sum([e[0].count("-") for e in alignment]))
            / min(len_seq_a, len_seq_b)
            * 100
            )

        if identity <= 40.0:
            log.warning(
                f"\"Structural\" identity of chain {chain} is {identity:.2f}%,"
                " please check the results carefully"
                )
        else:
            log.info(
                f"\"Structural\" identity of chain {chain} is {identity:.2f}%"
                )

        # logging.debug("Reading alignment and matching numbering")
        for element in alignment:
            line_a, line_b = element

            resnum_a, seq_a, _ = line_a.split()
            resnum_b, seq_b, _ = line_b.split()

            resnum_a = int(resnum_a) - 1
            resnum_b = int(resnum_b) - 1

            for resname_a, resname_b in zip(seq_a, seq_b):
                if resname_a != "-":
                    resnum_a += 1

                if resname_b != "-":
                    resnum_b += 1

                if resname_a != "-" and resname_b != "-":
                    numbering_dic[chain][resnum_b] = resnum_a

    izone_fname = Path(output_path, "lovoalign.izone")
    log.debug(f"Saving .izone to {izone_fname.name}")
    dump_as_izone(izone_fname, numbering_dic)

    return numbering_dic

def align_seq(reference, model, output_path):
    """Sequence align and get the numbering relationship."""
    seqdic_a = pdb2fastadic(reference)
    seqdic_b = pdb2fastadic(model)

    if seqdic_a.keys() != seqdic_b.keys():
        # TODO: Implement chain-matching here
        return False

    align_dic = {}
    for a, b in zip(seqdic_a, seqdic_b):

        align_dic[a] = {}

        seq_a = Seq("".join(seqdic_a[a].values()))
        seq_b = Seq("".join(seqdic_b[b].values()))

        aligner = Align.PairwiseAligner()
        aligner.substitution_matrix = substitution_matrices.load("BLOSUM62")
        alns = aligner.align(seq_a, seq_b)
        top_aln = alns[0]

        aln_fname = Path(output_path, f"blosum62_{a}.aln")
        log.debug(f"Writing alignment to {aln_fname.name}")
        with open(aln_fname, "w") as fh:
            fh.write(str(top_aln))
        aligned_seg_a, aligned_seg_b = top_aln.aligned

        # this should always be true
        assert len(aligned_seg_a) == len(aligned_seg_b)

        identity = (
            str(top_aln).count("|") / float(min(len(seq_a), len(seq_b)))
            ) * 100

        if not any(e for e in top_aln.aligned):
            # No alignment!
            log.warning(
                f"No alignment for chain {a} is it protein/dna? "
                "Matching sequentially"
                )
            if all("X" in s for s in seq_a) and all("X" in s for s in seq_b):
                # this sequence contains only ligands, do it manually
                if len(seq_a) != len(seq_b):
                    # we cannot handle this
                    raise f"Cannot align chain {b}"
                for res_a, res_b in zip(seqdic_a[a], seqdic_b[b]):
                    align_dic[a].update({res_a: res_b})
        else:
            if identity <= 40.0:
                # Identity is very low
                log.warning(
                    f"Sequence identity of chain {a} is {identity:.2f}%,"
                    " please check the results carefully"
                    )
                log.warning(
                    "Please use alignment_method = \"structure\" instead")
            else:
                log.info(f"Sequence identity of chain {a} is {identity:.2f}%")
            for seg_a, seg_b in zip(aligned_seg_a, aligned_seg_b):
                start_a, end_a = seg_a
                start_b, end_b = seg_b

                reslist_a = list(seqdic_a[a].keys())[start_a:end_a]
                reslist_b = list(seqdic_b[b].keys())[start_b:end_b]

                align_dic[a].update(
                    dict((i, j) for i, j in zip(reslist_a, reslist_b))
                    )
    izone_fname = Path(output_path, "blosum62.izone")
    log.debug(f"Saving .izone to {izone_fname.name}")
    dump_as_izone(izone_fname, align_dic)

    return align_dic

def make_range(chain_range_dic):
    """Expand a chain dictionary into ranges."""
    chain_ranges = {}
    for chain in chain_range_dic:
        min_idx = min(chain_range_dic[chain])
        max_idx = max(chain_range_dic[chain])
        chain_ranges[chain] = (min_idx, max_idx)
    return chain_ranges

def dump_as_izone(fname, numbering_dic):
    """Dump the numbering dictionary as .izone."""
    # FIXME: Collapse the izones so its faster to load in profit
    with open(fname, "w") as fh:
        for chain in numbering_dic:
            for bound_res in numbering_dic[chain]:
                unbound_res = numbering_dic[chain][bound_res]
                #
                izone_str = (
                    "ZONE "
                    f"{chain}{bound_res}:{chain}{unbound_res}"
                    f"{os.linesep}"
                    )
                fh.write(izone_str)

# # debug only
# def write_coord_dic(output_name, coord_dic):
#     """Add a dummy atom to a PDB file according to a list of coordinates."""
#     with open(output_name, "w") as fh:
#         for i, k in enumerate(coord_dic):
#             atom_num = f"{i+1}".rjust(4, " ")
#             chain, resnum, atom = k
#             resnum = int(resnum)
#             resnum = f"{resnum}".rjust(3, " ")
#             atom_name = f"{atom}".rjust(3, " ")
#             x, y, z = coord_dic[k]
#             dum_x = f"{x:.3f}".rjust(7, " ")
#             dum_y = f"{y:.3f}".rjust(7, " ")
#             dum_z = f"{z:.3f}".rjust(7, " ")
#             dummy_line = (
#                 f"ATOM   {atom_num} {atom_name}  DUM {chain} {resnum}   "
#                 f"  {dum_x} {dum_y} {dum_z}  1.00  1.00   "
#                 "        H  " + os.linesep
#                 )
#             fh.write(dummy_line)

# # debug only
# def write_coords(output_name, coor_list):
#     """Add a dummy atom to a PDB file according to a list of coordinates."""
#     with open(output_name, "w") as fh:
#         for i, dummy_coord in enumerate(coor_list):
#             atom_num = f"{i}".rjust(4, " ")
#             resnum = f"{i}".rjust(3, " ")
#             dum_x = f"{dummy_coord[0]:.3f}".rjust(7, " ")
#             dum_y = f"{dummy_coord[1]:.3f}".rjust(7, " ")
#             dum_z = f"{dummy_coord[2]:.3f}".rjust(7, " ")
#             dummy_line = (
#                 f"ATOM   {atom_num}  H   DUM X {resnum}   "
#                 f"  {dum_x} {dum_y} {dum_z}  1.00  1.00   "
#                 "        H  " + os.linesep
#                 )
#             fh.write(dummy_line)

# # debug only
# def write_pymol_viz(resdic):
#     """Write PyMol vizualitation."""
#     for k in resdic:
#         reslist = "+".join(map(str, resdic[k]))
#         cmd = f"sele {k}, chain {k} and resid {reslist}"
#         print(cmd)