from pandas.core.frame import DataFrame
import torch
import torch.nn as nn
import numpy as np
import time
import argparse
import random
from torch.multiprocessing import Pool

from loss import CustomLoss
from dist_loss import CustomDistLoss
from datagen import get_data_loader
from model import PIXOR
from dist_model import PIXOR_DIST
from dist_model2 import PIXOR_DIST2
from utils import get_model_name, load_config, get_logger, plot_bev, plot_label_map, plot_pr_curve, get_bev
from postprocess import filter_pred, compute_matches, compute_ap
import torch.nn.utils.prune as prune
import pandas as pd

### MAIN FILE OF PIXOR NEURAL NETWORK ###
### TO RUN AN EXPERIMENT, CALL ONE OF THE ARGUMENTS LISTED IN MAIN ###

# Constructs and returns the base PIXOR model, the corresponding loss function,
# optimizer, and learning rate scheduler
def build_model(config, device, train=True):
    net = PIXOR(config['geometry'], config['use_bn'])
    loss_fn = CustomLoss(device, config, num_classes=1)

    if torch.cuda.device_count() <= 1:
        config['mGPUs'] = False
    if config['mGPUs']:
        print("using multi gpu")
        net = nn.DataParallel(net)

    net = net.to(device)
    loss_fn = loss_fn.to(device)
    if not train:
        return net, loss_fn

    optimizer = torch.optim.SGD(net.parameters(
    ), lr=config['learning_rate'], momentum=config['momentum'], weight_decay=config['weight_decay'])
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=config['lr_decay_at'], gamma=0.1)

    return net, loss_fn, optimizer, scheduler

# Constructs the distilled PIXOR models, specifiying which one using [distill].
# Returns the model, corresponding loss function, optimizer, and learning rate scheduler
def build_dist_model(config, device, train=True, distill="distill_layer"):

    if distill == "distill_layer":
        net = PIXOR_DIST(config['geometry'], config['use_bn'])

    if distill == "distill_conv":
        net = PIXOR_DIST2(config['geometry'], config['use_bn'])

    loss_fn = CustomDistLoss(device, config, num_classes=1)

    if torch.cuda.device_count() <= 1:
        config['mGPUs'] = False
    if config['mGPUs']:
        print("using multi gpu")
        net = nn.DataParallel(net)  # parallelizing net

    net = net.to(device)
    loss_fn = loss_fn.to(device)

    if not train:
        return net, loss_fn

    optimizer = torch.optim.SGD(net.parameters(
    ), lr=config['learning_rate'], momentum=config['momentum'], weight_decay=config['weight_decay'])
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=config['lr_decay_at'], gamma=0.1)

    return net, loss_fn, optimizer, scheduler

# Evaluates a batch of data given in [loader] for distilled models. Needs the
# teacher model [teacher_net] and distilled loss [loss_fn].
# Returns various metrics and image predictions
def eval_batch_distill(config, net, teacher_net, loss_fn, loader, device, eval_range='all'):

    net.eval()
    teacher_net.eval()  # even though I'm pretty sure it's already in eval mode

    if config['mGPUs']:
        net.module.set_decode(True)
        teacher_net.module.set_decode(True)
    else:
        net.set_decode(True)
        teacher_net.set_decode(True)

    cls_loss = 0
    loc_loss = 0
    cls_teacher_loss = 0
    loc_teacher_loss = 0
    all_scores = []
    all_matches = []
    log_images = []
    gts = 0
    preds = 0
    t_fwd = 0
    t_nms = 0

    log_img_list = random.sample(range(len(loader.dataset)), 10)

    with torch.no_grad():
        for i, data in enumerate(loader):
            tic = time.time()
            input, label_map, image_id = data
            input = input.to(device)
            label_map = label_map.to(device)
            tac = time.time()
            predictions = net(input)
            t_fwd += time.time() - tac

            # getting the teacher predictions, no need to get forward times for this set
            teacher_predictions = teacher_net(input)

            loss, cls, loc, cls_teacher, loc_teacher = loss_fn(
                predictions, teacher_predictions, label_map)
            cls_loss += cls
            loc_loss += loc
            cls_teacher_loss += cls_teacher
            loc_teacher_loss += loc_teacher
            t_fwd += (time.time() - tic)

            toc = time.time()
            # Parallel post-processing
            predictions = list(torch.split(predictions.cpu(), 1, dim=0))
            batch_size = len(predictions)
            # with Pool(processes=3) as pool:
            #    preds_filtered = pool.starmap(
            #        filter_pred, [(config, pred) for pred in predictions])

            preds_filtered = []
            for pred in predictions:
                preds_filtered.append(filter_pred(config, pred))

            t_nms += (time.time() - toc)
            args = []

            for j in range(batch_size):
                _, label_list = loader.dataset.get_label(image_id[j].item())
                corners, scores = preds_filtered[j]
                gts += len(label_list)
                preds += len(scores)
                all_scores.extend(list(scores))
                if image_id[j] in log_img_list:
                    input_np = input[j].cpu().permute(1, 2, 0).numpy()
                    pred_image = get_bev(input_np, corners)
                    log_images.append(pred_image)

                arg = (np.array(label_list), corners, scores)
                args.append(arg)

            # Parallel compute matchesi

            with Pool(processes=3) as pool:
                matches = pool.starmap(compute_matches, args)

            for j in range(batch_size):
                all_matches.extend(list(matches[j][1]))

            # print(time.time() -tic)

    all_scores = np.array(all_scores)
    all_matches = np.array(all_matches)
    sort_ids = np.argsort(all_scores)
    all_matches = all_matches[sort_ids[::-1]]

    metrics = {}
    try:
        AP, precisions, recalls, precision, recall = compute_ap(
            all_matches, gts, preds)
    except:
        AP, precisions, recalls, precision, recall = [-1]*5

    metrics['AP'] = AP
    metrics['Precision'] = precision
    metrics['Recall'] = recall
    metrics['Forward Pass Time'] = t_fwd/len(loader.dataset)
    metrics['Postprocess Time'] = t_nms/len(loader.dataset)

    cls_loss = cls_loss / len(loader)
    loc_loss = loc_loss / len(loader)
    cls_teacher_loss = cls_teacher_loss / len(loader)
    loc_teacher_loss = loc_teacher_loss / len(loader)
    metrics['loss'] = cls_loss + loc_loss + cls_teacher_loss + loc_teacher_loss

    return metrics, precisions, recalls, log_images

# Evaluates a batch of data in [loader] for a default model. Returns various
# metrics and image predictions
def eval_batch(config, net, loss_fn, loader, device, eval_range='all'):
    net.eval()
    if config['mGPUs']:
        net.module.set_decode(True)
    else:
        net.set_decode(True)

    cls_loss = 0
    loc_loss = 0
    all_scores = []
    all_matches = []
    log_images = []
    gts = 0
    preds = 0
    t_fwd = 0
    t_nms = 0

    log_img_list = random.sample(range(len(loader.dataset)), 10)

    with torch.no_grad():
        for i, data in enumerate(loader):
            tic = time.time()
            input, label_map, image_id = data
            input = input.to(device)
            label_map = label_map.to(device)
            tac = time.time()
            predictions = net(input)
            t_fwd += time.time() - tac
            loss, cls, loc = loss_fn(predictions, label_map)
            cls_loss += cls
            loc_loss += loc
            t_fwd += (time.time() - tic)

            toc = time.time()
            # Parallel post-processing
            predictions = list(torch.split(predictions.cpu(), 1, dim=0))
            batch_size = len(predictions)
            # with Pool(processes=3) as pool:
            #    preds_filtered = pool.starmap(
            #        filter_pred, [(config, pred) for pred in predictions])

            preds_filtered = []
            for pred in predictions:
                preds_filtered.append(filter_pred(config, pred))

            t_nms += (time.time() - toc)
            args = []

            for j in range(batch_size):
                _, label_list = loader.dataset.get_label(image_id[j].item())
                corners, scores = preds_filtered[j]
                gts += len(label_list)
                preds += len(scores)
                all_scores.extend(list(scores))
                if image_id[j] in log_img_list:
                    input_np = input[j].cpu().permute(1, 2, 0).numpy()
                    pred_image = get_bev(input_np, corners)
                    log_images.append(pred_image)

                arg = (np.array(label_list), corners, scores)
                args.append(arg)

            # Parallel compute matchesi

            with Pool(processes=3) as pool:
                matches = pool.starmap(compute_matches, args)

            for j in range(batch_size):
                all_matches.extend(list(matches[j][1]))

            # print(time.time() -tic)

    all_scores = np.array(all_scores)
    all_matches = np.array(all_matches)
    sort_ids = np.argsort(all_scores)
    all_matches = all_matches[sort_ids[::-1]]

    metrics = {}
    try:
        AP, precisions, recalls, precision, recall = compute_ap(
            all_matches, gts, preds)
    except:
        AP, precisions, recalls, precision, recall = [-1]*5

    metrics['AP'] = AP
    metrics['Precision'] = precision
    metrics['Recall'] = recall
    metrics['Forward Pass Time'] = t_fwd/len(loader.dataset)
    metrics['Postprocess Time'] = t_nms/len(loader.dataset)

    cls_loss = cls_loss / len(loader)
    loc_loss = loc_loss / len(loader)
    metrics['loss'] = cls_loss + loc_loss

    return metrics, precisions, recalls, log_images

# Evaluates a set of data given by [loader]. Basically does the same thing
# as eval_batch, but doesn't include as many predictions
def eval_dataset(config, net, loss_fn, loader, device, e_range='all'):
    net.eval()
    if config['mGPUs']:
        net.module.set_decode(True)
    else:
        net.set_decode(True)

    t_fwds = 0
    t_post = 0
    loss_sum = 0

    img_list = range(len(loader.dataset))
    if e_range != 'all':
        e_range = min(e_range, len(loader.dataset))
        img_list = random.sample(img_list, e_range)

    log_img_list = random.sample(img_list, 10)

    gts = 0
    preds = 0
    all_scores = []
    all_matches = []
    log_images = []

    with torch.no_grad():
        for image_id in img_list:
            # tic = time.time()
            num_gt, num_pred, scores, pred_image, pred_match, loss, t_forward, t_nms = eval_one(net, loss_fn, config, loader,
                                                                                image_id, device, plot=False)
            gts += num_gt
            preds += num_pred
            loss_sum += loss
            all_scores.extend(list(scores))
            all_matches.extend(list(pred_match))

            t_fwds += t_forward
            t_post += t_nms

            if image_id in log_img_list:
                log_images.append(pred_image)
            # print(time.time() - tic)

    all_scores = np.array(all_scores)
    all_matches = np.array(all_matches)
    sort_ids = np.argsort(all_scores)
    all_matches = all_matches[sort_ids[::-1]]

    metrics = {}
    AP, precisions, recalls, precision, recall = compute_ap(
        all_matches, gts, preds)
    metrics['AP'] = AP
    metrics['Precision'] = precision
    metrics['Recall'] = recall
    metrics['loss'] = loss_sum / len(img_list)
    metrics['Forward Pass Time'] = t_fwds / len(img_list)
    metrics['Postprocess Time'] = t_post / len(img_list)

    return metrics, precisions, recalls, log_images

# Training loop for training a default model from scratch. This code
# was originally written by Phillip Huang in his implementation of PIXOR.
def train_orig(exp_name, device, amp=False):
    # Load Hyperparameters
    config, learning_rate, batch_size, max_epochs = load_config(exp_name)

    # Dataset and DataLoader
    train_data_loader, test_data_loader = get_data_loader(batch_size, config['use_npy'],
                                                          geometry=config['geometry'], frame_range=config['frame_range'])
    # Model
    net, loss_fn, optimizer, scheduler = build_model(
        config, device, train=True)

    # Tensorboard Logger
    train_logger = get_logger(config, 'train')
    val_logger = get_logger(config, 'val')

    if config['resume_training']:
        saved_ckpt_path = get_model_name(config)
        if config['mGPUs']:
            net.module.load_state_dict(torch.load(
                saved_ckpt_path, map_location=device))
        else:
            net.load_state_dict(torch.load(
                saved_ckpt_path, map_location=device))
        print("Successfully loaded trained ckpt at {}".format(saved_ckpt_path))
        st_epoch = config['resume_from']
    else:
        # writefile(config, 'train_loss.csv', 'iteration, cls_loss, loc_loss\n')
        # writefile(config, 'val_loss.csv', 'epoch, cls_loss, loc_loss\n')
        st_epoch = 0

    step = 1 + st_epoch * len(train_data_loader)
    cls_loss = 0
    loc_loss = 0
    scaler = torch.cuda.amp.GradScaler()

    best_loss = 10E8
    num_val_not_decreasing = 0
    for epoch in range(st_epoch, max_epochs):
        start_time = time.time()

        train_loss = 0

        net.train()
        if config['mGPUs']:
            net.module.set_decode(False)
        else:
            net.set_decode(False)
        scheduler.step()

        for input, label_map, image_id in train_data_loader:

            tic = time.time()  # print('step', step)
            input = input.to(device)
            label_map = label_map.to(device)
            optimizer.zero_grad()

            # Use Automatic Mixed Precision
            if (amp):
                with torch.cuda.amp.autocast():
                    predictions = net(input)
                    loss, cls, loc = loss_fn(predictions, label_map)
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
            # Use float32
            else:
                predictions = net(input)
                loss, cls, loc = loss_fn(predictions, label_map)
                loss.backward()
                optimizer.step()

            cls_loss += cls
            loc_loss += loc
            train_loss += loss.item()

            if step % config['log_every'] == 0:
                cls_loss = cls_loss / config['log_every']
                loc_loss = loc_loss / config['log_every']
                train_logger.scalar_summary('cls_loss', cls_loss, step)
                train_logger.scalar_summary('loc_loss', loc_loss, step)
                cls_loss = 0
                loc_loss = 0

                # for tag, value in net.named_parameters():
                #    tag = tag.replace('.', '/')
                #    train_logger.histo_summary(tag, value.data.cpu().numpy(), step)
                #    train_logger.histo_summary(tag + '/grad', value.grad.data.cpu().numpy(), step)

            step += 1
            # print(time.time() - tic)

        # Record Training Loss
        print('train_loss', train_loss)
        print('train_data_loader', train_data_loader)
        train_loss = train_loss / len(train_data_loader)
        train_logger.scalar_summary('loss', train_loss, epoch + 1)
        print("Epoch {}|Time {:.3f}|Training Loss: {:.5f}".format(
            epoch + 1, time.time() - start_time, train_loss))

        # Run Validation every 2 epochs
        if (epoch + 1) % 2 == 0:
            tic = time.time()
            val_metrics, _, _, log_images = eval_batch(
                config, net, loss_fn, test_data_loader, device)
            for tag, value in val_metrics.items():
                val_logger.scalar_summary(tag, value, epoch + 1)
            val_logger.image_summary('Predictions', log_images, epoch + 1)
            print("Epoch {}|Time {:.3f}|Validation Loss: {:.5f}".format(
                epoch + 1, time.time() - tic, val_metrics['loss']))

        # Save Checkpoint
        #if (epoch + 1) == (st_epoch + max_epochs) or (epoch + 1) % config['save_every'] == 0:
        if val_metrics["loss"] < best_loss:
            best_loss = val_metrics["loss"]
            model_path = get_model_name(config, exp_name = exp_name, epoch = epoch + 1)
            if config['mGPUs']:
                torch.save(net.module.state_dict(), model_path)
            else:
                torch.save(net.state_dict(), model_path)
            print("Checkpoint saved at {}".format(model_path))
            num_val_not_decreasing = 0
        else:
            num_val_not_decreasing+=1
            # early stopping
            if num_val_not_decreasing == 10:
                print(f"early stopping on epoch {str(epoch)} for experiement {exp_name}")
                break

    # torch.save(net.state_dict(), model_path)
    print('Finished Training')

# Training loop for a distilled modeled - requires a teacher model [teacher_net].
def train_distilled(net, loss_fn, optimizer, scheduler, teacher_net, device, config, learning_rate, batch_size, max_epochs, exp_name=None):

    df_logs = pd.DataFrame()

    # Dataset and DataLoader - using same data for distilled model (re training on same data)
    train_data_loader, test_data_loader = get_data_loader(
        batch_size, config['use_npy'], geometry=config['geometry'], frame_range=config['frame_range'])

    # Tensorboard Logger
    train_logger = get_logger(config, 'train')
    val_logger = get_logger(config, 'val')

    # setting the starting epoch
    if config['resume_training']:
        st_epoch = config['resume_from']
    else:
        st_epoch = 0

    # setting the initial step
    step = 1 + st_epoch * len(train_data_loader)
    # cross entropy loss - calculated adding the cross entropy loss wrt data and teacher model
    cls_loss = 0
    # regression loss - calculated adding regression loss wrt data and teacher model
    loc_loss = 0
    # cross entropy loss wrt to the teacher model
    cls_teacher_loss = 0
    # regression loss wrt to the teacher model
    loc_teacher_loss = 0

    best_loss = 100000000000
    num_val_not_decreasing = 0

    # training for a certain number of epochs
    for epoch in range(st_epoch, st_epoch + max_epochs):

        print("Epoch number: " + str(epoch + 1))

        start_time = time.time()  # starting timer

        train_loss = 0  # iteratively calculating training loss

        net.train()  # setting net to training mode (vs net.eval())

        # whether to set decode to false if using mGPUs
        if config['mGPUs']:
            net.module.set_decode(False)
            teacher_net.module.set_decode(False)
        else:
            net.set_decode(False)
            teacher_net.set_decode(False)

        for input, label_map, image_id, in train_data_loader:

            tic = time.time()  # print('step', step)
            input = input.to(device)
            label_map = label_map.to(device)
            optimizer.zero_grad()

            # Forward
            predictions = net(input)
            teacher_predictions = teacher_net(input)

            loss, cls, loc, cls_teacher, loc_teacher = loss_fn(
                predictions, teacher_predictions, label_map)

            loss.backward()
            optimizer.step()

            # keeping sum total of all the losses - will average later
            cls_loss += cls
            loc_loss += loc
            cls_teacher_loss += cls_teacher
            loc_teacher_loss += loc_teacher
            train_loss += loss.item()

            # logging cls_loss and loc_loss every certain number of steps
            if step % config['log_every'] == 0:
                cls_loss = cls_loss / config['log_every']
                loc_loss = loc_loss / config['log_every']
                cls_teacher_loss = cls_teacher_loss / config['log_every']
                loc_teacher_loss = loc_teacher_loss / config['log_every']
                train_logger.scalar_summary('cls_loss', cls_loss, step)
                train_logger.scalar_summary('loc_loss', loc_loss, step)
                train_logger.scalar_summary(
                    'cls_teacher_loss', cls_teacher_loss, step)
                train_logger.scalar_summary(
                    'loc_teacher_loss', loc_teacher_loss, step)
                cls_loss = 0
                loc_loss = 0
                cls_teacher_loss = 0
                loc_teacher_loss = 0

            # increasing # of steps
            step += 1

        scheduler.step()  # stepping the learning rate (after the first optimizer step)

        # Record Training Loss
        train_loss = train_loss / len(train_data_loader)
        train_logger.scalar_summary('loss', train_loss, epoch + 1)
        print("Epoch {}|Time {:.3f}|Training Loss: {:.5f}".format(
            epoch + 1, time.time() - start_time, train_loss))

        # Run Validation
        # if (epoch + 1) % 2 == 0:
        tic = time.time()
        val_metrics, _, _, log_images = eval_batch_distill(
            config, net, teacher_net, loss_fn, test_data_loader, device)
        for tag, value in val_metrics.items():
            val_logger.scalar_summary(tag, value, epoch + 1)
        val_logger.image_summary('Predictions', log_images, epoch + 1)
        print("Epoch {}|Time {:.3f}|Validation Loss: {:.5f}".format(
            epoch + 1, time.time() - tic, val_metrics['loss']))
        print(val_metrics)

        # Save Checkpoint, implemented with early stopping if validation loss is too high
        # if (epoch + 1) == (st_epoch + max_epochs) or (epoch + 1) % config['save_every'] == 0:
        if val_metrics["loss"] < best_loss:
            best_loss = val_metrics["loss"]
            model_path = get_model_name(
                config, exp_name=exp_name, epoch=epoch + 1)
            if config['mGPUs']:
                torch.save(net.module.state_dict(), model_path)
            else:
                torch.save(net.state_dict(), model_path)
            print("Checkpoint saved at {}".format(model_path))
            num_val_not_decreasing = 0
        else:
            num_val_not_decreasing += 1
            # early stopping
            if num_val_not_decreasing == 5:
                print(
                    f"early stopping on epoch {str(epoch)} for experiement {exp_name}")
                break

    df_logs = pd.concat(
        [
            df_logs,
            pd.DataFrame(
                {
                    "exp_name": exp_name,
                    "AP": val_metrics["AP"],
                    "Precision": val_metrics["Precision"],
                    "Recall": val_metrics["Recall"],
                    "Forward Pass Time": val_metrics["Forward Pass Time"],
                    "Postprocess Time": val_metrics["Postprocess Time"],
                    "loss": val_metrics["loss"],
                    "train_loss": train_loss,
                }, index=[1]
            ),
        ]
    )

    print('Finished Training')

    return df_logs

# Training loop for a pruned model
def train(net, device, config, learning_rate, batch_size, max_epochs, exp_name=None):

    df_logs = pd.DataFrame()

    # Dataset and DataLoader
    train_data_loader, test_data_loader = get_data_loader(batch_size, config['use_npy'],
                                                          geometry=config['geometry'], frame_range=config['frame_range'])

    # Tensorboard Logger
    train_logger = get_logger(config, exp_name=exp_name, mode='train')
    val_logger = get_logger(config, exp_name=exp_name, mode='val')

    if config['resume_training']:
        st_epoch = config['resume_from']
    else:
        # writefile(config, 'train_loss.csv', 'iteration, cls_loss, loc_loss\n')
        # writefile(config, 'val_loss.csv', 'epoch, cls_loss, loc_loss\n')
        st_epoch = 0

    step = 1 + st_epoch * len(train_data_loader)
    cls_loss = 0
    loc_loss = 0

    best_loss = 100000000000
    num_val_not_decreasing = 0

    for epoch in range(st_epoch, st_epoch + max_epochs):

        start_time = time.time()

        train_loss = 0

        net.train()
        if config['mGPUs']:
            net.module.set_decode(False)
        else:
            net.set_decode(False)
        scheduler.step()

        for input, label_map, image_id in train_data_loader:

            tic = time.time()  # print('step', step)
            input = input.to(device)
            label_map = label_map.to(device)
            optimizer.zero_grad()

            # Forward
            predictions = net(input)
            loss, cls, loc = loss_fn(predictions, label_map)
            loss.backward()
            optimizer.step()
            cls_loss += cls
            loc_loss += loc
            train_loss += loss.item()

            if step % config['log_every'] == 0:
                cls_loss = cls_loss / config['log_every']
                loc_loss = loc_loss / config['log_every']
                train_logger.scalar_summary('cls_loss', cls_loss, step)
                train_logger.scalar_summary('loc_loss', loc_loss, step)
                cls_loss = 0
                loc_loss = 0

                # for tag, value in net.named_parameters():
                #    tag = tag.replace('.', '/')
                #    train_logger.histo_summary(tag, value.data.cpu().numpy(), step)
                #    train_logger.histo_summary(tag + '/grad', value.grad.data.cpu().numpy(), step)

            step += batch_size
            # print(time.time() - tic)

        # Record Training Loss
        train_loss = train_loss / len(train_data_loader)
        train_logger.scalar_summary('loss', train_loss, epoch + 1)
        print("Epoch {}|Time {:.3f}|Training Loss: {:.5f}".format(
            epoch + 1, time.time() - start_time, train_loss))

        # Run Validation
        # if (epoch + 1) % 2 == 0:
        tic = time.time()
        val_metrics, _, _, log_images = eval_batch(
            config, net, loss_fn, test_data_loader, device)
        for tag, value in val_metrics.items():
            val_logger.scalar_summary(tag, value, epoch + 1)
        val_logger.image_summary('Predictions', log_images, epoch + 1)
        print("Epoch {}|Time {:.3f}|Validation Loss: {:.5f}".format(
            epoch + 1, time.time() - tic, val_metrics['loss']))
        print(val_metrics)

        # Save Checkpoint
        # if (epoch + 1) == (st_epoch + max_epochs) or (epoch + 1) % config['save_every'] == 0:
        if val_metrics["loss"] < best_loss:
            num_val_not_decreasing = 0
            best_loss = val_metrics["loss"]
            model_path = get_model_name(
                config, exp_name=exp_name, epoch=epoch + 1)
            if config['mGPUs']:
                torch.save(net.module.state_dict(), model_path)
            else:
                torch.save(net.state_dict(), model_path)
            print("Checkpoint saved at {}".format(model_path))
        else:
            num_val_not_decreasing += 1
            # early stopping
            if num_val_not_decreasing == 4:
                print(
                    f"early stopping on epoch {str(epoch)} for experiement {exp_name}")
                break

    df_logs = pd.concat(
        [
            df_logs,
            pd.DataFrame(
                {
                    "exp_name": exp_name,
                    "AP": val_metrics["AP"],
                    "Precision": val_metrics["Precision"],
                    "Recall": val_metrics["Recall"],
                    "Forward Pass Time": val_metrics["Forward Pass Time"],
                    "Postprocess Time": val_metrics["Postprocess Time"],
                    "loss": val_metrics["loss"],
                    "train_loss": train_loss,
                }, index=[1]
            ),
        ]
    )

    print('Finished Training')

    return df_logs

def eval_one(net, loss_fn, config, loader, image_id, device, plot=False, verbose=False):
    input, label_map, image_id = loader.dataset[image_id]
    input = input.to(device)
    label_map, label_list = loader.dataset.get_label(image_id)
    loader.dataset.reg_target_transform(label_map)
    label_map = torch.from_numpy(label_map).permute(
        2, 0, 1).unsqueeze_(0).to(device)

    # Forward Pass
    t_start = time.time()
    pred = net(input.unsqueeze(0))
    t_forward = time.time() - t_start

    loss, cls_loss, loc_loss = loss_fn(pred, label_map)
    pred.squeeze_(0)
    cls_pred = pred[0, ...]

    if verbose:
        print("Forward pass time", t_forward)

    # Filter Predictions
    t_start = time.time()
    corners, scores = filter_pred(config, pred)
    t_post = time.time() - t_start

    if verbose:
        print("Non max suppression time:", t_post)

    gt_boxes = np.array(label_list)
    gt_match, pred_match, overlaps = compute_matches(gt_boxes,
                                                     corners, scores, iou_threshold=0.5)

    num_gt = len(label_list)
    num_pred = len(scores)
    input_np = input.cpu().permute(1, 2, 0).numpy()
    pred_image = get_bev(input_np, corners)

    if plot == True:
        # Visualization
        plot_bev(input_np, label_list, window_name='GT')
        plot_bev(input_np, corners, window_name='Prediction')
        plot_label_map(cls_pred.numpy())

    return num_gt, num_pred, scores, pred_image, pred_match, loss.item(), t_forward, t_post

def experiment(exp_name, device, eval_range='all', plot=True):
    config, _, _, _ = load_config(exp_name)
    net, loss_fn = build_model(config, device, train=False)
    state_dict = torch.load(get_model_name(config), map_location=device)
    if config['mGPUs']:
        net.module.load_state_dict(state_dict)
    else:
        net.load_state_dict(state_dict)
    train_loader, val_loader = get_data_loader(config['batch_size'], config['use_npy'], geometry=config['geometry'],
                                               frame_range=config['frame_range'])

    # Train Set
    train_metrics, train_precisions, train_recalls, _ = eval_batch(
        config, net, loss_fn, train_loader, device, eval_range)
    print("Training mAP", train_metrics['AP'])
    fig_name = "PRCurve_train_" + config['name']
    legend = "AP={:.1%} @IOU=0.5".format(train_metrics['AP'])
    plot_pr_curve(train_precisions, train_recalls, legend, name=fig_name)

    # Val Set
    val_metrics, val_precisions, val_recalls, _ = eval_batch(
        config, net, loss_fn, val_loader, device, eval_range)

    print("Validation mAP", val_metrics['AP'])
    print("Net Fwd Pass Time on average {:.4f}s".format(
        val_metrics['Forward Pass Time']))
    print("Nms Time on average {:.4f}s".format(
        val_metrics['Postprocess Time']))

    fig_name = "PRCurve_val_" + config['name']
    legend = "AP={:.1%} @IOU=0.5".format(val_metrics['AP'])
    plot_pr_curve(val_precisions, val_recalls, legend, name=fig_name)

def test(exp_name, device, image_id):
    config, _, _, _ = load_config(exp_name)
    net, loss_fn = build_model(config, device, train=False)
    net.load_state_dict(torch.load(
        get_model_name(config), map_location=device))
    net.set_decode(True)
    train_loader, val_loader = get_data_loader(1, config['use_npy'], geometry=config['geometry'],
                                               frame_range=config['frame_range'])
    net.eval()

    with torch.no_grad():
        num_gt, num_pred, scores, pred_image, pred_match, loss, t_forward, t_nms = \
            eval_one(net, loss_fn, config, train_loader,
                     image_id, device, plot=True)

        TP = (pred_match != -1).sum()
        print("Loss: {:.4f}".format(loss))
        print("Precision: {:.2f}".format(TP/num_pred))
        print("Recall: {:.2f}".format(TP/num_gt))
        print("forward pass time {:.3f}s".format(t_forward))
        print("nms time {:.3f}s".format(t_nms))

# Prunes a model given a certain pruning method
def prune_model(model, method="L1_unstructured", prune_amt=0.2):

    params_to_prune = []

    for mod_name, nn_mod in model.named_modules():

        if isinstance(nn_mod, torch.nn.Conv2d):
            if nn_mod.bias is not None:
                params_to_prune.append((nn_mod, "bias"))

    if method == "L1_unstructured":
        prune.global_unstructured(
            params_to_prune,
            pruning_method=prune.L1Unstructured,
            amount=prune_amt,
        )
    elif method == "random_unstructured":
        prune.global_unstructured(
            params_to_prune,
            pruning_method=prune.RandomUnstructured,
            amount=prune_amt,
        )

# main function call
if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='PIXOR custom implementation')
    parser.add_argument(
        '--mode', choices=['train', 'val', 'test', 'prune-fine-tune', 'distill', 'hardware-performance', 'hardware-performance-distill', 'amp'], help='name of the experiment')
    parser.add_argument('--name', required=True, help="name of the experiment")
    parser.add_argument('--device', default='cpu', help='device to train on')
    parser.add_argument('--eval_range', type=int, help="range of evaluation")
    parser.add_argument('--test_id', type=int, default=25,
                        help="id of the image to test")
    parser.add_argument('--weights', type=str, help='path of weights to load')
    args = parser.parse_args()

    device = torch.device(args.device)
    if not torch.cuda.is_available():
        device = torch.device('cpu')
    print("Using device", device)

    if args.mode == 'train':
        train_orig(args.name, device)

        # Load Hyperparameters
        # config, learning_rate, batch_size, max_epochs = load_config(args.name)
        # train(device, config, learning_rate, batch_size, max_epochs)
    if args.mode == 'val':
        if args.eval_range is None:
            args.eval_range = 'all'
        experiment(args.name, device, eval_range=args.eval_range, plot=False)
    if args.mode == 'test':
        test(args.name, device, image_id=args.test_id)

    if args.mode == "distill":
        # loading model config
        config, learning_rate, batch_size, max_epochs = load_config(args.name)
        # loading teacher config
        default_config, _, _, _ = load_config('default')

        # building model
        net, loss_fn, optimizer, scheduler = build_dist_model(
            config, device, train=True, distill=args.name)

        # loading the teacher model ('34epoch')
        teacher_net, _ = build_model(default_config, device, train=False)
        teacher_net.load_state_dict(torch.load(
            "experiments/default/34epoch", map_location=device))

        teacher_net.eval()  # setting to evaluate mode

        df_logs = train_distilled(net, loss_fn, optimizer, scheduler, teacher_net, device, config,
                                  learning_rate, batch_size, max_epochs, exp_name=args.name)

        df_logs.to_csv(args.name + "_experiment_logs.csv")

    if args.mode == 'prune-fine-tune':

        # Load Hyperparameters
        config, learning_rate, batch_size, _ = load_config(args.name)
        config["resume_training"] = True
        config["max_epochs"] = 5

        df_logs = pd.DataFrame()
        for prune_method in ["random_unstructured", "L1_unstructured"]:
            # 2,4,8,16, and 32 pruning reductation as recommended in the state of pruning paper
            for prune_amt in [0, 0.5, 1-0.25, 1-0.125, 1-0.0625, 1-0.03125]:

                # Model
                net, loss_fn, optimizer, scheduler = build_model(
                    config, device, train=True)

                if config['resume_training']:
                    saved_ckpt_path = get_model_name(config)
                    net.load_state_dict(torch.load(
                        saved_ckpt_path, map_location=device))
                    print("Successfully loaded trained ckpt at {}".format(
                        saved_ckpt_path))

                prune_model(net, method=prune_method, prune_amt=prune_amt)

                df_logs_temp = train(net, device, config, learning_rate, batch_size,
                                     config["max_epochs"], exp_name=prune_method + "_" + str(prune_amt))

                df_logs = pd.concat([df_logs, df_logs_temp])

        df_logs.to_csv("prune_experiment_logs.csv")

    if args.mode == 'hardware-performance-distill':
        config, _, _, _ = load_config(args.name)
        net, _, _, _ = build_dist_model(
            config, device, train=True, distill=args.name)

        weights_path = args.weights

        net.load_state_dict(torch.load(weights_path, map_location=device))
        print("Successfully loaded trained ckpt at {}".format(weights_path))

        train_data_loader, _ = get_data_loader(config['batch_size'], config['use_npy'],
                                               geometry=config['geometry'], frame_range=config['frame_range'])

        inference_time = 0
        total_frames = 0

        net.eval()
        net.set_decode(False)

        loading_and_inference_time_start = time.time()

        for input, _, _ in train_data_loader:

            total_frames += len(input)

            tic = time.time()  # print('step', step)
            input = input.to(device)

            inference_time_start = time.time()
            predictions = net(input)
            inference_time += time.time() - inference_time_start

        total_loading_and_inference_time = time.time() - loading_and_inference_time_start

        average_loading_plus_inference = 1.0 * \
            total_loading_and_inference_time / total_frames
        average_inference = 1.0*inference_time / total_frames

        print(
            f"average time to load and predict on an image: {average_loading_plus_inference}")
        print(f"average time to predict on an image: {average_inference}")

    if args.mode == 'hardware-performance':

        config, _, _, _ = load_config(args.name)

        net, _, _, _ = build_model(config, device, train=True)

        weights_path = args.weights

        saved_ckpt_path = get_model_name(config)
        net.load_state_dict(torch.load(saved_ckpt_path, map_location=device))
        print("Successfully loaded trained ckpt at {}".format(saved_ckpt_path))

        prune_model(net)

        net.load_state_dict(torch.load(
            weights_path, map_location=device))
        print("Successfully loaded trained ckpt at {}".format(weights_path))

        for mod_name, nn_mod in net.named_modules():

            if isinstance(nn_mod, torch.nn.Conv2d):
                if nn_mod.bias is not None:
                    nn_mod = torch.nn.utils.prune.remove(nn_mod, name='bias')

        train_data_loader, _ = get_data_loader(config['batch_size'], config['use_npy'],
                                               geometry=config['geometry'], frame_range=config['frame_range'])

        inference_time = 0
        total_frames = 0

        net.eval()
        net.set_decode(False)

        loading_and_inference_time_start = time.time()

        for input, _, _ in train_data_loader:

            total_frames += len(input)

            tic = time.time()  # print('step', step)
            input = input.to(device)

            inference_time_start = time.time()
            predictions = net(input)
            inference_time += time.time() - inference_time_start

        total_loading_and_inference_time = time.time() - loading_and_inference_time_start

        average_loading_plus_inference = 1.0 * \
            total_loading_and_inference_time / total_frames
        average_inference = 1.0*inference_time / total_frames

        print(
            f"average time to load and predict on an image: {average_loading_plus_inference}")
        print(f"average time to predict on an image: {average_inference}")
    # Automatic Mixed Precision
    if args.mode == 'amp':
        train_orig(args.name, device, amp=True)
    # before launching the program! CUDA_VISIBLE_DEVICES=0, 1 python main.py .......


import csv
import re
import glob
import numpy as np
from keras.models import load_model
import face_recognition as fr

def main():
    model = load_model("face_rank_model.h5")
    files = glob.glob(r"data/*")
    for f in files:
        image = fr.load_image_file(f)
        encs = fr.face_encodings(image)
        if len(encs) != 1:
            print("Find %d faces in %s" % (len(encs), f))
            continue
        predicted = model.predict(np.array(encs))
        predicted = np.squeeze(predicted)
        list1=[re.split('\/',f)[-1],predicted]
        with open("result.csv", "a", newline='') as fff:
            csv.writer(fff).writerow(list1)
        print(type(predicted))
        print("%s: %.4f" % (f, predicted))

if __name__ == '__main__':
    main()


import math
import pdb
import sys
import time
import numpy as np

import torch
import torchvision.models.detection.mask_rcnn
from .utils import *
from PIL import Image
from PIL import ImageDraw
from pathlib import Path
from torchvision import transforms as T
from .coco_eval import CocoEvaluator
from .coco_utils import get_coco_api_from_dataset

data_dir = Path.home() / 'hdrive' / 'data'
wgisdPath = data_dir / 'wgisd'

def draw_boxes(img, boxes, fname, save_dir):
   assert fname is not None
   draw = ImageDraw.Draw(img)
   for box in boxes:
      box = list(np.array(box))
      draw.rectangle(box)
   path = wgisdPath / save_dir
   if not path.exists():
      path.mkdir(parents=True)
   img.save(path / fname)

def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq,
                    update_weights=True):
    model.train()
    metric_logger = MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", SmoothedValue(window_size=1, fmt="{value:.6f}"))
    header = f"Epoch: [{epoch}]"

    lr_scheduler = None
    if update_weights is True and epoch == 0:
        warmup_factor = 1.0 / 1000
        warmup_iters = min(1000, len(data_loader) - 1)

        lr_scheduler = torch.optim.lr_scheduler.LinearLR(
            optimizer, start_factor=warmup_factor, total_iters=warmup_iters
        )

    loss_list = []
    for images, targets in metric_logger.log_every(data_loader, print_freq, header):
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)

        losses = sum(loss for loss in loss_dict.values())

        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = reduce_dict(loss_dict)
        losses_reduced = sum(loss for loss in loss_dict_reduced.values())

        loss_value = losses_reduced.item()
        loss_list.append(loss_value)

        if not math.isfinite(loss_value):
            print(f"Loss is {loss_value}, stopping training")
            print(loss_dict_reduced)
            sys.exit(1)

        if update_weights is True:
           optimizer.zero_grad()
           losses.backward()
           optimizer.step()

        if lr_scheduler is not None:
           lr_scheduler.step()

        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])

    return metric_logger, loss_list

def _get_iou_types(model):
    model_without_ddp = model
    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
        model_without_ddp = model.module
    iou_types = ["bbox"]
    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
        iou_types.append("segm")
    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
        iou_types.append("keypoints")
    return iou_types

@torch.inference_mode()
def evaluate(model, data_loader, device, save_results=False, save_dir='img_boxes'):
    model.to(device)
    dataset = data_loader.dataset
    n_threads = torch.get_num_threads()
    # FIXME remove this and make paste_masks_in_image run on the GPU
    torch.set_num_threads(1)
    cpu_device = torch.device("cpu")
    model.eval()
    metric_logger = MetricLogger(delimiter="  ")
    header = "Test:"

    coco = get_coco_api_from_dataset(data_loader.dataset)
    iou_types = _get_iou_types(model)
    coco_evaluator = CocoEvaluator(coco, iou_types)

    for images, targets in metric_logger.log_every(data_loader, 100, header):
        images = list(img.to(device) for img in images)

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        model_time = time.time()
        outputs = model(images)

        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
        model_time = time.time() - model_time

        res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
        evaluator_time = time.time()
        coco_evaluator.update(res)
        evaluator_time = time.time() - evaluator_time
        metric_logger.update(model_time=model_time, evaluator_time=evaluator_time)

        if save_results:
           save_dir = Path(save_dir)
           for img, target, output in zip(images, targets, outputs):
              fname = dataset.get_fname(target['image_id'])
              gtImg = T.ToPILImage()(img)
              resultImg = T.ToPILImage()(img)
              draw_boxes(gtImg, target['boxes'], fname, save_dir/'gt')
              draw_boxes(resultImg, output['boxes'].detach().numpy(),
                         fname, save_dir/'results')

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    coco_evaluator.synchronize_between_processes()

    # accumulate predictions from all images
    coco_evaluator.accumulate()
    coco_evaluator.summarize()
    torch.set_num_threads(n_threads)
    return coco_evaluator


# coding=utf-8
# Copyright 2019 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility for parsing artifacts about a R2R house specification."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import csv
import json
import math
import os
import re
import numpy as np

import tensorflow.compat.v2 as tf

from valan.r2r import house_utils

_BANNED_MP40_CAT_INDEX = {0, 40, 41}
# Note: `oriented_bbox` is another namedtuple of ('axis0', 'axis1', 'radii')
RoomObject = collections.namedtuple('RoomObject', [
    'center', 'distance', 'name', 'category', 'clean_category', 'oriented_bbox'
])

class Region(object):
  """Region data in the house segmentation file."""

  NUM_TOKENS = 20  # Number of tokens to expect in line.
  INDEX_LOC = 1    # Location of region index.
  LEVEL_LOC = 2    # Location of level index.
  LABEL_LOC = 5    # Location of label.
  PX_LOC = 6       # Location of px.
  PY_LOC = 7       # Location of py.
  PZ_LOC = 8       # Location of pz.

  def __init__(self, region_line):
    parts = region_line.strip().split()
    assert 'R' == parts[0]
    assert self.NUM_TOKENS == len(parts)
    self.index = int(parts[self.INDEX_LOC])
    self.level_index = int(parts[self.LEVEL_LOC])
    self.label = parts[self.LABEL_LOC]
    self.center = (float(parts[self.PX_LOC]), float(parts[self.PY_LOC]),
                   float(parts[self.PZ_LOC]))

class Pano(object):
  """Pano data in the house segmentation file."""

  NUM_TOKENS = 13  # Number of tokens to expect in line.
  NAME_LOC = 1     # Location of name.
  INDEX_LOC = 2    # Location of index.
  REGION_LOC = 3   # Location of region index.
  PX_LOC = 5       # Location of px.
  PY_LOC = 6       # Location of py.
  PZ_LOC = 7       # Location of pz.

  def __init__(self, pano_line):
    parts = pano_line.strip().split()
    assert 'P' == parts[0]
    assert self.NUM_TOKENS == len(parts)
    self.name = parts[self.NAME_LOC]
    self.index = int(parts[self.INDEX_LOC])
    self.region_index = int(parts[self.REGION_LOC])
    self.center = (float(parts[self.PX_LOC]), float(parts[self.PY_LOC]),
                   float(parts[self.PZ_LOC]))
    # List of images (`Image`) of the pano.
    self.images = []

  def get_available_headings(self):
    # Only need to consider the heading angles of one of the 3 cameras.
    return [image.heading for image in self.images if image.camera_index == 0]

  def get_images_at(self, heading, threshold=math.pi / 12):
    """Returns the images within threshold arc distance from given heading."""
    angle_distances = [
        house_utils.compute_arc_distance(heading, image.heading)
        for image in self.images
    ]
    matches = [
        self.images[idx]
        for idx, dist in enumerate(angle_distances)
        if dist <= threshold
    ]
    # Sort matches by pitch angle for consistency.
    sorted_images = sorted(matches, key=lambda image: image.pitch)
    return sorted_images

class Image(object):
  """Image data in the house segmentation file."""

  NUM_TOKENS = 41                  # Number of tokens to expect in line.
  INDEX_LOC = 1                    # Location of image index.
  PANO_INDEX_LOC = 2               # Location of panorama index.
  NAME_LOC = 3                     # Location of the name of the image.
  CAMERA_INDEX_LOC = 4             # Location of the camera index.
  HEADING_INDEX_LOC = 5            # Location of the heading (yaw) index.
  EXTRINSIC_MATRIX_START_LOC = 6   # Start location of extrinsic camera matrix.
  EXTRINSIC_MATRIX_END_LOC = 18    # End location of extrinsic camera matrix.

  def __init__(self, image_line, scan_id):
    self.scan_id = scan_id
    parts = image_line.strip().split()
    assert 'I' == parts[0]
    assert self.NUM_TOKENS == len(parts)
    self.index = int(parts[self.INDEX_LOC])
    self.pano_index = int(parts[self.PANO_INDEX_LOC])
    self.name = parts[self.NAME_LOC]
    self.camera_index = int(parts[self.CAMERA_INDEX_LOC])
    self.heading_index = int(parts[self.HEADING_INDEX_LOC])

    # Compute heading and pitch from extrinsing matrix.
    extrinsic_coordinates = [
        float(coord) for coord in
        parts[self.EXTRINSIC_MATRIX_START_LOC:self.EXTRINSIC_MATRIX_END_LOC]
    ]
    rotation_matrix = [
        extrinsic_coordinates[i:i + 3]
        for i in range(0, len(extrinsic_coordinates), 4)
    ]
    self.heading, self.pitch, _ = house_utils.get_euler_angles(rotation_matrix)

class Category(object):
  """Category data in the house segmentation file."""

  NUM_TOKENS = 11            # Number of tokens to expect in line.
  INDEX_LOC = 1              # Location of index.
  CAT_MAP_INDEX_LOC = 2      # Location of category mapping index.
  CAT_MAP_NAME_LOC = 3       # Location of category mapping name.
  MPCAT_INDEX_LOC = 4        # Location of mpcat index.
  MPCAT_NAME_LOC = 5         # Location of mpcat name.

  def __init__(self, category_line, category_map=None):
    """Extract object index and category name for a line in the .house file."""
    parts = category_line.strip().split()
    assert 'C' == parts[0]
    assert self.NUM_TOKENS == len(parts)
    self.index = int(parts[self.INDEX_LOC])
    self.category_mapping_index = int(parts[self.CAT_MAP_INDEX_LOC])
    # Raw category name
    self.category_mapping_name = ' '.join(
        parts[self.CAT_MAP_NAME_LOC].split('#'))
    self.mpcat40_index = int(parts[self.MPCAT_INDEX_LOC])
    self.mpcat40_name = parts[self.MPCAT_NAME_LOC]
    # Cleaned category name
    if category_map:
      self.clean_category_name = self._get_clean_cat_name(
          self.category_mapping_index, category_map)

  def _get_clean_cat_name(self,
                          category_index,
                          category_map,
                          count_cutoff_threshold=5):
    """Map category index to a clean category name instead of raw categeory.

    The clean categories are from the R2R `category_mapping.tsv` file. It
    corrects typos and standardizes the raw categories, which is much more fine
    grained than the mpcat40 categories (which only has 40 categories).
    For more information see:
    https://github.com/niessner/Matterport/blob/master/metadata/category_mapping.tsv

    Args:
      category_index: int; the category mapping index extracted from the
        category line from the .house file.
      category_map: a dict returned by `_load_cat_map()`, containing mappings
        from category index to category names, mpcat40 name, and count.
      count_cutoff_threshold: categories with counts below the threshold are
        replaced with their corresponding mpcat40 names. This is used to
        truncate the long tail of rarely used category names.

    Returns:
      A unicode string for the clean category name.
    """
    cat_map = category_map[category_index]
    if cat_map['count'] >= count_cutoff_threshold:
      clean_name = cat_map['clean_category']
    else:
      clean_name = cat_map['mpcat40']
    return clean_name

class Object(object):
  """Object data in the house segmentation file."""

  NUM_TOKENS = 24  # Number of tokens to expect in line.
  INDEX_LOC = 1    # Location of index.
  REGION_LOC = 2   # Location of index.
  CAT_LOC = 3      # Location of index.
  PX_LOC = 4       # Location of px.
  PY_LOC = 5       # Location of py.
  PZ_LOC = 6       # Location of pz.
  AXIS0_X_LOC = 7    # Location of X axis min of the oriented bbox
  AXIS0_Y_LOC = 8    # Location of Y axis min of the oriented bbox
  AXIS0_Z_LOC = 9    # Location of Z axis min of the oriented bbox
  AXIS1_X_LOC = 10   # Location of X axis max of the oriented bbox
  AXIS1_Y_LOC = 11   # Location of Y axis max of the oriented bbox
  AXIS1_Z_LOC = 12   # Location of Z axis max of the oriented bbox
  RADIUS_X_LOC = 13  # Location of X raidus of the oriented bbox
  RADIUS_Y_LOC = 14  # Location of Y raidus of the oriented bbox
  RADIUS_Z_LOC = 15  # Location of Z raidus of the oriented bbox

  def __init__(self, object_line):
    parts = object_line.strip().split()
    assert 'O' == parts[0]
    assert self.NUM_TOKENS == len(parts)
    self.index = int(parts[self.INDEX_LOC])
    self.region_index = int(parts[self.REGION_LOC])
    self.category_index = int(parts[self.CAT_LOC])
    self.center = (float(parts[self.PX_LOC]), float(parts[self.PY_LOC]),
                   float(parts[self.PZ_LOC]))
    # Oriented bounding box (obbox)
    oriented_bbox = collections.namedtuple('OrientedBbox',
                                           ['axis0', 'axis1', 'radii'])
    axis0 = (float(parts[self.AXIS0_X_LOC]), float(parts[self.AXIS0_Y_LOC]),
             float(parts[self.AXIS0_Z_LOC]))
    axis1 = (float(parts[self.AXIS1_X_LOC]), float(parts[self.AXIS1_Y_LOC]),
             float(parts[self.AXIS1_Z_LOC]))
    radii = (float(parts[self.RADIUS_X_LOC]), float(parts[self.RADIUS_Y_LOC]),
             float(parts[self.RADIUS_Z_LOC]))
    self.obbox = oriented_bbox(axis0, axis1, radii)

  def is_well_formed(self):
    return (self.index >= 0 and self.region_index >= 0 and
            self.category_index >= 0)

class R2RHouseParser(object):
  """Parser to extract various annotations in a house to assist perception."""

  def __init__(self,
               house_file_path,
               category_map_dir=None,
               category_map_file='category_mapping.tsv',
               banned_mp40_cat_index=None):
    """Parses regions, panos, categories and objects from house spec file.

    For more information see:
    https://github.com/niessner/Matterport/blob/master/data_organization.md

    Args:
      house_file_path: Path to scan id house specification file.
      category_map_dir: Dir of category mapping file 'category_mapping.tsv'.
        If not provided, then the `clean_category` will be omitted.
      category_map_file: str; file name for the category mapping table. Use
        default unless the mapping table is different.
      banned_mp40_cat_index: A set of mpcat40 category indices, e.g., {0, 41}
        for (void, unlabled). If provided, then these categories will be ignored
        when extracting objects of each pano.
    """
    # Load category map and banned mp40 categories.
    if category_map_dir:
      assert tf.io.gfile.isdir(
          category_map_dir), '{} must be an existing dir.'.format(
              category_map_dir)
      category_map = _load_cat_map(
          category_map_file, category_map_dir, delimiter='\t')
    else:
      # Default to None and omit `clean_category_name` if dir is not given.
      category_map = None

    if not banned_mp40_cat_index:
      self.banned_mp40_cat_index = _BANNED_MP40_CAT_INDEX
    else:
      self.banned_mp40_cat_index = banned_mp40_cat_index

    self.scan_id = os.path.splitext(os.path.basename(house_file_path))[0]
    with tf.io.gfile.GFile(house_file_path, 'r') as input_file:
      # Skip but check header line.
      assert re.match('^ASCII .*', input_file.readline().strip()) is not None
      house_info = input_file.readline().strip().split()
      assert 29 == len(house_info)
      self.num_images = int(house_info[3])
      self.num_panos = int(house_info[4])
      self.num_objects = int(house_info[8])
      self.num_categories = int(house_info[9])
      self.num_regions = int(house_info[10])
      self.regions = {}
      self.panos = {}
      self.categories = {}
      self.objects = {}
      self.images = {}
      for line in input_file:
        if line[0] == 'R':
          r = Region(line)
          assert r.index not in self.regions
          self.regions[r.index] = r
        elif line[0] == 'P':
          p = Pano(line)
          assert p.index not in self.panos
          self.panos[p.index] = p
        elif line[0] == 'C':
          c = Category(line, category_map)
          assert c.index not in self.categories
          self.categories[c.index] = c
        elif line[0] == 'O':
          o = Object(line)
          assert o.index not in self.objects
          self.objects[o.index] = o
        elif line[0] == 'I':
          i = Image(line, self.scan_id)
          assert i.index not in self.images
          self.images[i.index] = i

      assert self.num_regions == len(self.regions)
      assert self.num_panos == len(self.panos)
      assert self.num_categories == len(self.categories)
      assert self.num_objects == len(self.objects)
      assert self.num_images == len(self.images)

    # Organize objects into regions for easy retrieval later.
    self.pano_name_map = {}
    for p in self.panos.values():
      self.pano_name_map[p.name] = p.index
    self.region_object_map = collections.defaultdict(list)
    for o in self.objects.values():
      self.region_object_map[o.region_index] += [o.index]

    # Add images to the associated panos.
    for image in self.images.values():
      pano = self.get_pano_by_name(image.name)
      pano.images.append(image)

  def __repr__(self):
    return 'Regions: {}, Panos: {}, Cats: {}, Objs: {}'.format(
        self.num_regions, self.num_panos, self.num_categories, self.num_objects)

  def get_pano_objects(self, pano_name):
    """Extract the set of objects given a pano.

    Only returns the closest object of the same mp40 category and skips any
    objects with mp40 category in `self.banned_mp40_cat_index` (e.g. misc,
    void, unlabeled categories).

    Args:
      pano_name: panoromic hash id.

    Returns:
      Dictionary where key is the center of the `RoomObject` and the
      value is the `RoomObject` named tuple. In particular,
      `RoomObject.oriented_bbox` is another named tuple (axis0, axis1, radii)
      containing the axis orientation of the bounding box and the radii of the
      object along each axis.
    """
    pano_id = self.pano_name_map.get(pano_name, None)
    if pano_id is None:  # Note that `pano_id` is int and thus can be 0.

      return {}

    room_objects = {}
    region_index = self.panos[pano_id].region_index
    pano_center = self.panos[pano_id].center
    for object_index in self.region_object_map[region_index]:
      try:
        category = self.categories[self.objects[object_index].category_index]
        # NOTE: 'unknown' objects are sometimes labeled as mpcat40=40 (misc)
        # instead of mpcat40=41 (unlabeled). So we specifically exclude it here.
        if ('unknown' not in category.category_mapping_name.lower()) and (
            category.mpcat40_index not in self.banned_mp40_cat_index):
          object_center = self.objects[object_index].center
          assert object_center not in room_objects, self.objects[object_index]
          room_objects[object_center] = RoomObject(
              object_center,
              np.linalg.norm(np.array(object_center) - np.array(pano_center)),
              category.category_mapping_name,
              category.mpcat40_name,
              (category.clean_category_name if hasattr(
                  category, 'clean_category_name') else None),
              self.objects[object_index].obbox)
      except KeyError:
        # Note that this happens because some objects have been marked with -1
        # categories. We can safely ignore these objects.
        assert self.objects[object_index].category_index == -1
    return room_objects

  def get_pano_by_name(self, pano_name):
    return self.panos[self.pano_name_map[pano_name]]

  def get_panos_graph(self,
                      connections_file,
                      cache_best_paths_and_distances=False):
    """Returns a house_utils.Graph object with the panos as nodes.

    The connectivity file should be a json with a single line, containing one
    dictionary for every pano in the house scan. The fields of this dictionary
    used in this method are 'image_id', a unique string that identifies the
    pano, and 'unobstructed', a list of booleans with length equal to the number
    of panos in the scan, representing whether there is an unobstructed direct
    path from the current pano to the pano with that index.

    A more detailed description of the format of the file can be found at
    https://github.com/peteanderson80/Matterport3DSimulator/tree/master/connectivity.

    Args:
      connections_file: Path to the file containing the connections.
      cache_best_paths_and_distances: Whether or not to cache the best path and
        distance for every pair of nodes in the graph.
    """
    assert tf.io.gfile.exists(connections_file), ('Missing required file: %s' %
                                                  connections_file)

    with tf.io.gfile.GFile(connections_file, 'r') as f:
      connections_info = json.loads(f.readline())

    pano_idx_map = {}  # Dictionary mapping idx in connection file to pano name.
    excluded_pano_names = set()  # Set of names of excluded panos.
    for idx, pano_info in enumerate(connections_info):
      pano_name = pano_info['image_id']
      pano_idx_map[idx] = pano_name
      if not pano_info['included']:
        excluded_pano_names.add(pano_name)

    # We build a dictionary indexed by the id of the panos, with values
    # being a dictionary of connected panos to a house_utils.ConnectionInfo
    # storing the distance and heading_angle between them.
    nodes_dict = collections.defaultdict(house_utils.NodeInfo)
    for pano_info in connections_info:
      current_pano = self.get_pano_by_name(pano_info['image_id'])
      if current_pano.name in excluded_pano_names:
        continue  # Don't add excluded panos to the graph.

      for index, is_unobstructed in enumerate(pano_info['unobstructed']):
        connected_pano = self.get_pano_by_name(pano_idx_map[index])
        # Path should be unobstructed and target point should be included in the
        # simulator.
        if is_unobstructed and connected_pano.name not in excluded_pano_names:
          distance = house_utils.compute_distance(current_pano.center,
                                                  connected_pano.center)
          heading_angle = house_utils.compute_heading_angle(
              current_pano.center,
              connected_pano.center,
              radians=True,
              apply_r2r_correction=True)
          pitch_angle = house_utils.compute_pitch_angle(
              current_pano.center, connected_pano.center, radians=True)

          connection_info = house_utils.ConnectionInfo(distance, heading_angle,
                                                       pitch_angle)
          nodes_dict[current_pano.name].connections[connected_pano.name] = (
              connection_info)

    graph = house_utils.Graph(nodes_dict, cache_best_paths_and_distances)
    for node in graph.nodes.keys():
      graph.nodes[node].coords = self.get_pano_by_name(node).center

    return graph

def _load_cat_map(category_map_file='category_mapping.tsv',
                  file_dir='',
                  delimiter='\t'):
  """Load category mapping table from file.

  The mapping table is available at:
  https://github.com/niessner/Matterport/blob/master/metadata/category_mapping.tsv

  Args:
    category_map_file: str; category mapping table file.
    file_dir: str; the dir to `category_map_file`.
    delimiter: str; optional delimiter for the input table.

  Returns:
    A dict that maps category index to category names and counts.
  """
  filepath = os.path.join(file_dir, category_map_file)
  assert tf.io.gfile.exists(filepath), (
      'Missing category mapping file: {}'.format(filepath))
  data = {}
  with tf.io.gfile.GFile(filepath, 'r') as f:
    reader = csv.reader(f, delimiter=delimiter)
    for i, row in enumerate(reader):
      assert len(row) == 18, 'Num columns must be 18.'
      if i == 0:
        header = row
        assert header[0:4] == ['index', 'raw_category', 'category', 'count']
        assert header[-2:] == ['mpcat40index', 'mpcat40']
      else:
        entries = [r.lower() for r in row]
        # entries[0] is the index of each line.
        data[int(entries[0])] = {
            'raw_category': entries[1],
            'clean_category':  # Only take the first part if has '\' in name.
                entries[2].split('/')[0].strip(),
            'count': int(entries[3]),
            'mpcat40index': int(entries[-2]),
            'mpcat40': entries[-1],
        }
  return data


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import socket
import time
import json
import numpy as np
import tensorflow as tf
import model_blstm
import atexit
import sys
import select
from tagProcess import transferS

sys.path.append('../../ui_labeling/preprocessing')
import sphere_fitting

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string('data_dir', '../data/',
                           "data directory")
tf.app.flags.DEFINE_string('checkpoints_dir', '../checkpoints/',
                           "training checkpoints directory")
tf.app.flags.DEFINE_string('log_dir', '../train_log/',
                           "summary directory")
tf.app.flags.DEFINE_string('restore_path', None,
                           "path of saving model eg: ../checkpoints/model.ckpt-5")
tf.app.flags.DEFINE_integer('batch_size', 128,
                            "mini-batch size")
tf.app.flags.DEFINE_integer('total_epoches', 300,
                            "total training epoches")
tf.app.flags.DEFINE_integer('hidden_size', 128,
                            "size of LSTM hidden memory")
tf.app.flags.DEFINE_integer('num_layers', 2,
                            "number of stacked blstm")
tf.app.flags.DEFINE_integer("input_dims", 10,
                            "input dimensions")
tf.app.flags.DEFINE_integer("num_classes", 69,  # 68 letters + 1 blank
                            "num_labels + 1(blank)")
tf.app.flags.DEFINE_integer('save_freq', 250,
                            "frequency of saving model")
tf.app.flags.DEFINE_float('learning_rate', 0.001,
                          "learning rate of RMSPropOptimizer")
tf.app.flags.DEFINE_float('decay_rate', 0.99,
                          "decay rate of RMSPropOptimizer")
tf.app.flags.DEFINE_float('momentum', 0.9,
                          "momentum of RMSPropOptimizer")
tf.app.flags.DEFINE_float('max_length', 1940,
                          "pad to same length")
tf.app.flags.DEFINE_integer('label_pad', 63,
                            "label pad size")

letter_table = [' ', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
                'g', 'ga', 'h', 'i', 'j', 'k', 'km', 'l', 'm', 'n', 'o', 'p', 'pt', 'q', 'r', 's', 'sc', 'sp', 't', 'u', 'v', 'w', 'x', 'y', 'z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '<b>']

class ModelConfig(object):
    """
    testing config
    """

    def __init__(self):
        self.data_dir = FLAGS.data_dir
        self.checkpoints_dir = FLAGS.checkpoints_dir
        self.log_dir = FLAGS.log_dir
        self.restore_path = FLAGS.restore_path
        self.batch_size = FLAGS.batch_size
        self.total_epoches = FLAGS.total_epoches
        self.hidden_size = FLAGS.hidden_size
        self.num_layers = FLAGS.num_layers
        self.input_dims = FLAGS.input_dims
        self.num_classes = FLAGS.num_classes
        self.save_freq = FLAGS.save_freq
        self.learning_rate = FLAGS.learning_rate
        self.decay_rate = FLAGS.decay_rate
        self.momentum = FLAGS.momentum
        self.max_length = FLAGS.max_length
        self.label_pad = FLAGS.label_pad

    def show(self):
        print("data_dir:", self.data_dir)
        print("checkpoints_dir:", self.checkpoints_dir)
        print("log_dir:", self.log_dir)
        print("restore_path:", self.restore_path)
        print("batch_size:", self.batch_size)
        print("total_epoches:", self.total_epoches)
        print("hidden_size:", self.hidden_size)
        print("num_layers:", self.num_layers)
        print("input_dims:", self.input_dims)
        print("num_classes:", self.num_classes)
        print("save_freq:", self.save_freq)
        print("learning_rate:", self.learning_rate)
        print("decay_rate:", self.decay_rate)
        print("momentum:", self.momentum)
        print("max_length:", self.max_length)
        print("label_pad:", self.label_pad)

server_address = ('140.113.210.19',2001)
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

sock.bind(server_address)
def whenexit():
    connection.close()
    print("Exit and close connection")

sock.listen(1)

# def recvall(sock):
#     BUFF_SIZE = 1  # 4 KiB
#     data = ""
#     while True:
#         part = sock.recv(1024)
#         print("One time")
#         print("Size ", sys.getsizeof(part))
#         if sys.getsizeof(part) < BUFF_SIZE:
#             data += str(part)
#             print("go here ")
#             # either 0 or end of data
#             break
#         data += str(part)
#     print("See!")
#     print(data)
#     return data

# def recvall(sock):
#     buf = sock.recv(4096)
#     while buf:
#         yield buf
#         print("one time ",time.time())
#         buf = sock.recv(4096)

def recvall(sock):
    framents = []
    while True:
        buf = sock.recv(1024)
        sentence = buf.decode('ascii')
        # print(sentence)
        mark = sentence[-3:]
        # print("mark ",mark)
        if sentence[-3:] == 'bye':
            framents.append(sentence[:-3])
            break
            # sock.sendall("good".encode('utf-8'))
        else:
            framents.append(sentence)
    return framents

# def recvall(sock):
#     data = sock.recv(5000000000)
#     return str(data)

with tf.get_default_graph().as_default() as graph:
    config = ModelConfig()
    config.show()
    model = model_blstm.HWRModel(config, graph)
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(init)
        saver.restore(sess, FLAGS.restore_path)

        print("restore")
        # exit()
        while True:
            connection, clientAddr = sock.accept()
            atexit.register(whenexit)
            try:
                print('connection ', clientAddr)
                while True:

                    data = ''.join(recvall(connection))
                    print(time.time())
                    # print(type(data))
                    # print(data[:30])
                    json_data = json.loads(data)
                    #print(sphere_fitting.vr_sphere_fitting(json_data))
                    input_data = transferS(sphere_fitting.vr_sphere_fitting(json_data))

                    # time.sleep(5)
                    # connection.sendall("Done".encode('utf-8'))
                    # connection.close()
                    # break
                    # process data

                    #################
                    seq_len_list = []
                    for _, v in enumerate(input_data):
                        seq_len_list.append(v.shape[0])
                    seq_len_list = np.array(seq_len_list).astype(np.int32)
                    padded_input_data = []
                    for _, v in enumerate(input_data):
                        residual = FLAGS.max_length - v.shape[0]
                        padding_array = np.zeros([residual, FLAGS.input_dims])
                        padded_input_data.append(
                            np.concatenate([v, padding_array], axis=0))
                    padded_input_data = np.array(padded_input_data)
                    num_batch = int(
                        padded_input_data.shape[0] / config.batch_size)
                    for i, _ in enumerate(padded_input_data):
                        start_time = time.time()
                        predict = model.predict(
                            sess, padded_input_data[i:i + 1], seq_len_list[i:i + 1])
                        str_decoded = ''.join(
                            [letter_table[x] for x in np.asarray(predict.values)])
                        end_time = time.time()
                        # print('Original val: %s' % label_data[i])
                        print('Decoded  val: %s' % str_decoded)
                        print('Time Cost: %f' % (end_time - start_time))
                        connection.sendall(str_decoded.encode('utf-8'))
                        connection.close()
                        print("inner connection close")
                    break
                    # print("ok?")
                    # print(type(data))

            finally:
                print("connection close")
                connection.close()


import gym
import numpy as np
from brain import Agent

env = gym.make('CartPole-v0')
N = 20
batch_size = 5
n_epochs = 4
alpha = 0.0003
agent = Agent(n_actions=env.action_space.n, batch_size=batch_size, alpha=alpha,
                n_epochs=n_epochs, input_dims=env.observation_space.shape)
agent.load_models()

observation = env.reset()
done = False
score = 0
while not done:
    action, prob, val = agent.choose_action(observation)
    observation_, reward, done, info = env.step(action)
    env.render()
    score += reward
    observation = observation_
print("Score:", score)


import numpy as np
import random

def get_response(self, question):
    # This part classifies the response as either "glossary", "general" or "events"
    user_bag = self.bag_of_words(question, self.wordbank)  # Creates a bag of words out of the question.
    predicted_results = self.model.predict([user_bag])  # Returns an np.array of probabilities for each of the tags in the wordbank.
    result_index = np.argmax(predicted_results)  # Finds the largest
    result_tag = self.tags[result_index]
    responses = []
    for pattern in self.dataset["intents"]:  # Finds the response that corresponds with the tag.
        if pattern['tag'] == result_tag:
            if predicted_results[0, result_index] > 0.3:  # If it is more than 30 percent sure that that is the correct response, return the response
                responses = pattern['response']
            else:
                responses = ["Sorry I did not quite understand that. Please try again!"]

    return random.choice(responses)


"""
Send multiple jobs to the cluster.
"""
import os
import subprocess
import time

import numpy as np

ALGOS = ["sac"]
ENVS = ["HalfCheetahBulletEnv-v0"]
N_SEEDS = 5
N_EVAL_EPISODES = 10
LOG_STD_INIT = [-6, -5, -4, -3, -2, -1, 0, 1]

os.makedirs(os.path.join("logs", "slurm"), exist_ok=True)

for algo in ALGOS:
    for env_id in ENVS:
        for log_std_init in LOG_STD_INIT:
            log_folder = f"logs_std_{np.exp(log_std_init):.4f}"
            for _ in range(N_SEEDS):
                args = [
                    "--algo",
                    algo,
                    "--env",
                    env_id,
                    "--hyperparams",
                    f'policy_kwargs:"dict(log_std_init={log_std_init}, net_arch=[400, 300])"',
                    "--eval-episodes",
                    N_EVAL_EPISODES,
                    "-f",
                    log_folder,
                    "-uuid",
                ]
                args = list(map(str, args))

                command = " ".join(["python", "-u", "train.py"] + args)

                ok = subprocess.call(["sbatch", "cluster_torchy.sh", algo, env_id, "ablation", command])
                time.sleep(0.05)


#!/usr/bin/env python3

import argparse
import copy
import cv2
import skvideo.io               # pip3 install scikit-video
import math
import navpy
import numpy as np
import os
import pandas as pd
import re
from scipy import interpolate  # strait up linear interpolation, nothing fancy
from tqdm import tqdm

from aurauas_flightdata import flight_loader, flight_interp

import camera
import correction
import correlate
from feat_data import FeatureData
import features
import hud
import hud_glass

# helpful constants
d2r = math.pi / 180.0
r2d = 180.0 / math.pi

# default sizes of primatives
render_w = 1920
render_h = 1080

# configure
experimental_overlay = False

parser = argparse.ArgumentParser(description='correlate video data to flight data.')
parser.add_argument('--flight', required=True, help='load specified aura flight log')
parser.add_argument('--video', required=True, help='original video')
parser.add_argument('--camera', help='select camera calibration file')
parser.add_argument('--cam-mount', choices=['forward', 'down', 'rear'],
                    default='forward',
                    help='approximate camera mounting orientation')
parser.add_argument('--rot180', action='store_true')
parser.add_argument('--scale', type=float, default=1.0, help='scale input')
parser.add_argument('--scale-preview', type=float, default=0.5,
                    help='scale preview')
parser.add_argument('--alpha', type=float, default=0.7, help='hud alpha blend')
parser.add_argument('--resample-hz', type=float, default=60.0,
                    help='resample rate (hz)')
parser.add_argument('--start-time', type=float, help='fast forward to this flight log time before begining video render.')
parser.add_argument('--time-shift', type=float, help='skip autocorrelation and use this offset time')
parser.add_argument('--plot', action='store_true', help='Plot stuff at the end of the run')
parser.add_argument('--auto-switch', choices=['old', 'new', 'none', 'on'], default='new', help='auto/manual switch logic helper')
parser.add_argument('--airspeed-units', choices=['kt', 'mps'], default='kt', help='display units for airspeed')
parser.add_argument('--altitude-units', choices=['ft', 'm'], default='ft', help='display units for airspeed')
parser.add_argument('--aileron-scale', type=float, default=1.0, help='useful for reversing aileron in display')
parser.add_argument('--elevator-scale', type=float, default=1.0, help='useful for reversing elevator in display')
parser.add_argument('--rudder-scale', type=float, default=1.0, help='useful for reversing rudder in display')
parser.add_argument('--flight-track-seconds', type=float, default=300.0, help='how many seconds of flight track to draw')
parser.add_argument('--keep-tmp-video', action='store_true', help='Keep the temp video')
parser.add_argument('--correction', help='correction table')
parser.add_argument('--features', help='feature database')
args = parser.parse_args()

smooth_cutoff_hz = 10

# for convenience
hz = args.resample_hz
r2d = 180.0 / math.pi
d2r = math.pi / 180.0

counter = 0

# pathname work
abspath = os.path.abspath(args.video)
filename, ext = os.path.splitext(abspath)
dirname = os.path.dirname(args.video)
video_rates = filename + "_rates.csv"
ekf_error = filename + "_error.csv"
local_config = dirname + "/camera.json"

# combinations that seem to work on linux
# ext = avi, fourcc = MJPG
# ext = avi, fourcc = XVID
# ext = m4v (was mov), fourcc = MP4V

ext = 'mp4'
tmp_video = filename + "_tmp." + ext
output_video = filename + "_hud." + ext

camera = camera.VirtualCamera()
camera.load(args.camera, local_config, args.scale)
cam_yaw, cam_pitch, cam_roll = camera.get_ypr()
K = camera.get_K()
dist = camera.get_dist()
print('Camera:', camera.get_name())
print('K:\n', K)
print('dist:', dist)

if os.path.exists(ekf_error):
    correction.load_horiz(ekf_error)

data, flight_format = flight_loader.load(args.flight)
print("imu records:", len(data['imu']))
print("gps records:", len(data['gps']))
if 'air' in data:
    print("airdata records:", len(data['air']))
print("filter records:", len(data['filter']))
if 'pilot' in data:
    print("pilot records:", len(data['pilot']))
if 'act' in data:
    print("act records:", len(data['act']))
if len(data['imu']) == 0 and len(data['gps']) == 0:
    print("not enough data loaded to continue.")
    quit()

interp = flight_interp.InterpolationGroup(data)
iter = flight_interp.IterateGroup(data)

# imu gyro data
imu = pd.DataFrame(data['imu'])
imu.set_index('time', inplace=True, drop=False)
imu_min = imu['time'].iat[0]
imu_max = imu['time'].iat[-1]
imu_count = len(imu)
imu_fs =  int(round((imu_count / (imu_max - imu_min))))
print("imu fs:", imu_fs)

# resample imu data
print("flight range = %.3f - %.3f (%.3f)" % (imu_min, imu_max,
                                             imu_max-imu_min))
flight_interp = []
flight_len = imu_max - imu_min
p_interp = interpolate.interp1d(imu['time'], imu['p'], bounds_error=False, fill_value=0.0)
q_interp = interpolate.interp1d(imu['time'], imu['q'], bounds_error=False, fill_value=0.0)
r_interp = interpolate.interp1d(imu['time'], imu['r'], bounds_error=False, fill_value=0.0)
alt_interp = interp.group['filter'].interp['alt']

for x in np.linspace(imu_min, imu_max, int(round(flight_len*hz))):
    flight_interp.append( [x, p_interp(x), q_interp(x), r_interp(x) ] )
print("flight len:", len(flight_interp))

# load camera rotation rate data (derived from feature matching video
# frames)
feat_data = FeatureData()
feat_data.load(video_rates)
feat_data.smooth(smooth_cutoff_hz)
feat_data.make_interp()
if args.plot:
    feat_data.plot()
feat_interp = feat_data.resample(args.resample_hz)

# find the time correlation of video vs flight data
time_shift = \
    correlate.sync_gyros(flight_interp, feat_interp, feat_data.span_sec,
                         hz=hz, cam_mount=args.cam_mount,
                         force_time_shift=args.time_shift, plot=args.plot)

# quick estimate ground elevation
sum = 0.0
count = 0
for f in data['filter']:
    air = interp.query(f['time'], 'air')
    if air['airspeed'] < 5.0:
        sum += f['alt']
        count += 1
if count > 0:
    ground_m = sum / float(count)
else:
    ground_m = data['filter'][0].alt
print("ground est:", ground_m)

# overlay hud(s)
hud1 = hud_glass.HUD()
hud2 = hud.HUD(K)

#cam_ypr = [-3.0, -12.0, -3.0] # yaw, pitch, roll
#ref = [44.7260320000, -93.0771072000, 0]
ref = [ data['gps'][0]['lat'], data['gps'][0]['lon'], 0.0 ]
hud1.set_ned_ref(data['gps'][0]['lat'], data['gps'][0]['lon'])
hud2.set_ned_ref(data['gps'][0]['lat'], data['gps'][0]['lon'])
print('ned ref:', ref)

print('temporarily disabling airport loading')
#hud1.load_airports()

hud1.set_ground_m(ground_m)
hud2.set_ground_m(ground_m)

if args.features:
    feats = features.load(args.features, ref)
    hud1.update_features(feats)
else:
    feats = []

metadata = skvideo.io.ffprobe(args.video)
#print(metadata.keys())
#import json
#print(json.dumps(metadata["video"], indent=4))
fps_string = metadata['video']['@avg_frame_rate']
(num, den) = fps_string.split('/')
fps = float(num) / float(den)
codec = metadata['video']['@codec_long_name']
w = int(round(int(metadata['video']['@width']) * args.scale))
h = int(round(int(metadata['video']['@height']) * args.scale))
total_frames = int(round(float(metadata['video']['@duration']) * fps))
print('fps:', fps)
print('codec:', codec)
print('output size:', w, 'x', h)

hud1.set_render_size(w, h)
hud2.set_render_size(w, h)

print("Opening ", args.video)
reader = skvideo.io.FFmpegReader(args.video, inputdict={}, outputdict={})

inputdict = {
    '-r': str(fps)
}

sane = {
    # See all options: https://trac.ffmpeg.org/wiki/Encode/H.264
    '-vcodec': 'libx264',  # use the h.264 codec
    '-pix_fmt': 'yuv420p', # support 'dumb' players
    '-crf': '17',          # visually lossless (or nearly so)
    '-preset': 'medium',   # default compression
    '-r': str(fps)         # fps
}
lossless = {
    # See all options: https://trac.ffmpeg.org/wiki/Encode/H.264
    '-vcodec': 'libx264',  # use the h.264 codec
    '-pix_fmt': 'yuv420p', # support 'dumb' players
    '-crf': '0',           # set the constant rate factor to 0, (lossless)
    '-preset': 'veryslow', # maximum compression
    '-r': str(fps)         # fps
}

writer = skvideo.io.FFmpegWriter(tmp_video, inputdict=inputdict, outputdict=sane)

last_time = 0.0

# set primative sizes based on rendered resolution.
size = math.sqrt(h*h + w*w)
hud1.set_line_width( int(round(size/1000.0)) )
hud1.set_font_size( size / 1400.0 )
hud1.set_color( hud.green2 )
hud1.set_units( args.airspeed_units, args.altitude_units)

hud2.set_line_width( int(round(size/1000.0)) )
hud2.set_font_size( size / 1400.0 )
hud2.set_color( hud.red2 )
hud2.set_units( args.airspeed_units, args.altitude_units)

filt_alt = None

if True and time_shift > 0:
    # catch up the flight path history (in case the video starts
    # mid-flight.)  Note: flight_min is the starting time of the filter data
    # set.
    print('seeding flight track ...')
    for time in np.arange(imu_min, time_shift, 1.0 / float(fps)):
        filt = interp.query(time, 'filter')
        #air = interp.query(time, 'air')
        gps = interp.query(time, 'gps')
        lat_deg = filt['lat']*r2d
        lon_deg = filt['lon']*r2d
        #altitude_m = air['alt_true']
        altitude_m = filt['alt']
        ned = navpy.lla2ned( lat_deg, lon_deg, altitude_m,
                             ref[0], ref[1], ref[2] )
        hud1.update_time(time, gps['unix_sec'])
        hud1.update_ned(ned, args.flight_track_seconds)
        hud1.update_events(data['event'])

shift_mod_hack = False
pbar = tqdm(total=int(total_frames), smoothing=0.05)
for frame in reader.nextFrame():
    frame = frame[:,:,::-1]     # convert from RGB to BGR (to make opencv happy)
    if args.rot180:
        frame = np.rot90(frame)
        frame = np.rot90(frame)

    time = float(counter) / fps + time_shift
    #print("frame: ", counter, "%.3f" % time, 'time shift:', time_shift)

    counter += 1
    if args.start_time and time < args.start_time:
        continue
    filt = interp.query(time, 'filter')
    gps = interp.query(time, 'gps')
    air = interp.query(time, 'air')
    ap = interp.query(time, 'ap')
    pilot = interp.query(time, 'pilot')
    act = interp.query(time, 'act')
    vn = filt['vn']
    ve = filt['ve']
    vd = filt['vd']
    psix = filt['psix']
    psiy = filt['psiy']
    yaw_rad = math.atan2(psiy, psix)
    pitch_rad = filt['the']
    roll_rad = filt['phi']
    if not correction.yaw_interp is None:
        yaw_rad += correction.yaw_interp(time)
    if not correction.pitch_interp is None:
        pitch_rad += correction.pitch_interp(time)
    if not correction.roll_interp is None:
        roll_rad += correction.roll_interp(time)
    lat_deg = filt['lat']*r2d
    lon_deg = filt['lon']*r2d
    #altitude_m = air['alt_true']
    altitude_m = filt['alt']
    if filt_alt == None:
        filt_alt = altitude_m
    else:
        filt_alt = 0.95 * filt_alt + 0.05 * altitude_m
    if 'airspeed' in air:
        airspeed_kt = air['airspeed']
    else:
        airspeed_kt = 0.0
    if 'wind_dir' in air:
        wind_deg = air['wind_dir']
        wind_kt = air['wind_speed']
    if 'alpha' in air and 'beta' in air:
        alpha_rad = air['alpha']*d2r
        beta_rad = air['beta']*d2r
        #print alpha_rad, beta_rad
    else:
        alpha_rad = None
        beta_rad = None
        #print 'no alpha/beta'
    if 'hdgx' in ap:
        ap_hdgx = ap['hdgx']
        ap_hdgy = ap['hdgy']
        ap_hdg = math.atan2(ap_hdgy, ap_hdgx)*r2d
        ap_roll = ap['roll']
        if not correction.roll_interp is None:
            ap_roll += correction.roll_interp(time) * r2d
        ap_pitch = ap['pitch']
        if not correction.pitch_interp is None:
            ap_pitch += correction.pitch_interp(time) * r2d
        ap_speed = ap['speed']
        ap_alt_ft = ap['alt']
    if 'aileron' in pilot:
        pilot_ail = pilot['aileron'] * args.aileron_scale
        pilot_ele = pilot['elevator'] * args.elevator_scale
        pilot_thr = pilot['throttle']
        pilot_rud = pilot['rudder'] * args.rudder_scale
        auto_switch = pilot['auto_manual']
    else:
        auto_switch = 0
    if 'aileron' in act:
        act_ail = act['aileron'] * args.aileron_scale
        act_ele = act['elevator'] * args.elevator_scale
        act_thr = act['throttle']
        act_rud = act['rudder'] * args.rudder_scale

    if args.auto_switch == 'none':
        flight_mode = 'manual'
    elif (args.auto_switch == 'new' and auto_switch < 0) or (args.auto_switch == 'old' and auto_switch > 0):
        flight_mode = 'manual'
    elif args.auto_switch == 'on':
        flight_mode = 'auto'
    else:
        flight_mode = 'auto'

    if 'mission' in data:
        excite_mode = mission['excite']
        test_index = mission['test_index']

    record = iter.next()
    hud1.update_task(record)
    while 'imu' in record and record['imu']['time'] < time:
        record = iter.next()
        hud1.update_task(record)

    ned = navpy.lla2ned( lat_deg, lon_deg, filt_alt,
                         ref[0], ref[1], ref[2] )
    #print 'ned:', ned
    camera.update_PROJ(ned, yaw_rad, pitch_rad, roll_rad)

    method = cv2.INTER_AREA
    #method = cv2.INTER_LANCZOS4
    frame_scale = cv2.resize(frame, (0,0), fx=args.scale, fy=args.scale,
                             interpolation=method)
    frame_undist = cv2.undistort(frame_scale, K, np.array(dist))

    # Create hud draw space
    if not experimental_overlay:
        hud1_frame = frame_undist.copy()
    else:
        hud1_frame = np.zeros((frame_undist.shape), np.uint8)

    hud1.update_time(time, gps['unix_sec'])
    if 'event' in data:
        hud1.update_events(data['event'])
    if 'mission' in data:
        hud1.update_test_index(excite_mode, test_index)
    hud1.update_camera(camera)
    hud1.update_cam_att(cam_yaw, cam_pitch, cam_roll)
    hud1.update_ned(ned, args.flight_track_seconds)
    hud1.update_lla([lat_deg, lon_deg, altitude_m])
    hud1.update_vel(vn, ve, vd)
    hud1.update_att_rad(roll_rad, pitch_rad, yaw_rad)
    if 'wind_dir' in air:
        hud1.update_airdata(airspeed_kt, altitude_m, wind_deg, wind_kt, alpha_rad, beta_rad)
    else:
        hud1.update_airdata(airspeed_kt, altitude_m)
    if 'hdgx' in ap:
        hud1.update_ap(flight_mode, ap_roll, ap_pitch, ap_hdg,
                       ap_speed, ap_alt_ft)
    else:
        hud1.update_ap(flight_mode, 0.0, 0.0, 0.0, 0.0, 0.0)
    if 'aileron' in pilot:
        hud1.update_pilot(pilot_ail, pilot_ele, pilot_thr, pilot_rud)
    if 'aileron' in act:
        hud1.update_act(act_ail, act_ele, act_thr, act_rud)
    if time >= imu_min and time <= imu_max:
        # only draw hud for time range when we have actual flight data
        hud1.update_frame(hud1_frame)
        hud1.draw()

    if not experimental_overlay:
        # weighted add of the HUD frame with the original frame to
        # emulate alpha blending
        alpha = args.alpha
        if alpha < 0: alpha = 0
        if alpha > 1: alpha = 1
        cv2.addWeighted(hud1_frame, alpha, frame_undist, 1 - alpha, 0, hud1_frame)
    else:
        # Now create a mask of hud and create its inverse mask also
        tmp = cv2.cvtColor(hud1_frame, cv2.COLOR_BGR2GRAY)
        ret, mask = cv2.threshold(tmp, 10, 255, cv2.THRESH_BINARY)
        mask_inv = cv2.bitwise_not(mask)

        # Now black-out the hud from the original image
        tmp_bg = cv2.bitwise_and(frame_undist, frame_undist, mask=mask_inv)

        # Put hud onto the main image
        hud1_frame = cv2.add(tmp_bg, hud1_frame)

    # cv2.imshow('hud', hud1_frame)
    cv2.imshow('hud', cv2.resize(hud1_frame, None, fx=args.scale_preview, fy=args.scale_preview))
    writer.writeFrame(hud1_frame[:,:,::-1])  #write the frame as RGB not BGR

    pbar.update(1)

    key = cv2.waitKeyEx(5)
    if key == -1:
        # no key press
        continue

    print('key:', key)

    if key == 27:
        break
    elif key == ord('y'):
        if shift_mod_hack:
            cam_yaw -= 0.5
        else:
            cam_yaw += 0.5
        camera.set_yaw(cam_yaw)
        camera.save(local_config)
        shift_mod_hack = False
    elif key == ord('p'):
        if shift_mod_hack:
            cam_pitch -= 0.5
        else:
            cam_pitch += 0.5
        camera.set_pitch(cam_pitch)
        camera.save(local_config)
        shift_mod_hack = False
    elif key == ord('r'):
        if shift_mod_hack:
            cam_roll += 0.5
        else:
            cam_roll -= 0.5
        camera.set_roll(cam_roll)
        camera.save(local_config)
        shift_mod_hack = False
    elif key == ord('-'):
        time_shift -= 1.0/60.0
        shift_mod_hack = False
    elif key == ord('+'):
        time_shift += 1.0/60.0
        shift_mod_hack = False
    elif key == 65505 or key == 65506:
        shift_mod_hack = True
pbar.close()

writer.close()
cv2.destroyAllWindows()

# now run ffmpeg as an external command to combine original audio
# track with new overlay video

# ex: ffmpeg -i opencv.avi -i orig.mov -c copy -map 0:v -map 1:a final.avi

from subprocess import call
result = call(["ffmpeg", "-an", "-i", tmp_video, "-vn", "-i", args.video, "-c:v", "copy", "-c:a", "aac", "-y", output_video])
print("ffmpeg result code:", result)
if result == 0 and not args.keep_tmp_video:
    print("removing temp video:", tmp_video)
    os.remove(tmp_video)
    print("output video:", output_video)