#!/usr/bin/env python
import argparse
import os
from tqdm import tqdm
import pprint as pp
import numpy as np
import h5py
import json
import torch
import torch.optim as optim
import torch.autograd as autograd
from torch.optim import lr_scheduler
from torch.autograd import Variable
from torch.utils.data import DataLoader
from tensorboard_logger import configure, log_value

from neural_combinatorial_rl.neural_combinatorial_rl import NeuralCombOptRL
from neural_combinatorial_rl.matching_nco import MatchingNeuralCombOptRL, MatchingNoDecoder
from envs import dataset

def str2bool(v):
      return v.lower() in ('true', '1')

parser = argparse.ArgumentParser(description="Neural Combinatorial Optimization with RL")

# Data
parser.add_argument('--task', default='sort_10-20', help="The task to solve, in the form {COP}_{size}, e.g., tsp_20")
parser.add_argument('--parallel_envs', type=int, default=128, help='')
parser.add_argument('--batch_size', type=int, default=128, help='')
parser.add_argument('--train_size', type=int, default=1000000, help='')
parser.add_argument('--val_size', type=int, default=10000, help='')
# Network
parser.add_argument('--embedding_dim', type=int, default=128, help='Dimension of input embedding')
parser.add_argument('--hidden_dim', type=int, default=128, help='Dimension of hidden layers in Enc/Dec')
parser.add_argument('--input_size', type=int, default=10)
parser.add_argument('--n_features', type=int, default=1)
parser.add_argument('--n_process_blocks', type=int, default=3, help='Number of process block iters to run in the Critic network')
parser.add_argument('--n_glimpses', type=int, default=2, help='No. of glimpses to use in the pointer network')
parser.add_argument('--use_tanh', type=str2bool, default=True)
parser.add_argument('--tanh_exploration', type=int, default=10, help='Hyperparam controlling exploration in the pointer net by scaling the tanh in the softmax')
parser.add_argument('--dropout', default=0., help='')
parser.add_argument('--terminating_symbol', default='<0>', help='')
parser.add_argument('--beam_size', default=1, help='Beam width for beam search')
# Training
parser.add_argument('--use_decoder', type=str2bool, default=False)
parser.add_argument('--actor_net_lr', default=1e-4, help="Set the learning rate for the actor network")
parser.add_argument('--critic_net_lr', default=1e-4, help="Set the learning rate for the critic network")
parser.add_argument('--actor_lr_decay_step', default=5000, help='')
parser.add_argument('--critic_lr_decay_step', default=5000, help='')
parser.add_argument('--actor_lr_decay_rate', default=0.96, help='')
parser.add_argument('--critic_lr_decay_rate', default=0.96, help='')
parser.add_argument('--reward_scale', default=2, type=float,  help='')
parser.add_argument('--is_train', type=str2bool, default=True, help='')
parser.add_argument('--n_epochs', default=1, help='')
parser.add_argument('--random_seed', type=int, default=24601, help='')
parser.add_argument('--max_grad_norm', default=1.0, help='Gradient clipping')
parser.add_argument('--use_cuda', type=str2bool, default=True, help='')
parser.add_argument('--critic_beta', type=float, default=0.9, help='Exp mvg average decay')
parser.add_argument('--use_KT', type=str2bool, default=True)
# Misc
parser.add_argument('--log_step', default=50, help='Log info every log_step steps')
parser.add_argument('--log_dir', type=str, default='results/logs')
parser.add_argument('--run_name', type=str, default='0')
parser.add_argument('--base_dir', type=str, default='/data/pemami/spg/')
parser.add_argument('--output_dir', type=str, default='outputs')
parser.add_argument('--epoch_start', type=int, default=0, help='Restart at epoch #')
parser.add_argument('--load_path', type=str, default='')
parser.add_argument('--disable_tensorboard', type=str2bool, default=True)
parser.add_argument('--plot_attention', type=str2bool, default=False)
parser.add_argument('--disable_progress_bar', type=str2bool, default=False)
parser.add_argument('--save_stats', type=str2bool, default=False)
parser.add_argument('--save_model', type=str2bool, default=False)
parser.add_argument('--_id', type=str, default='1234567')
parser.add_argument('--sl', type=str2bool, default=False)
parser.add_argument('--use_graph', type=str2bool, default=False)
parser.add_argument('--make_only', type=int, default=3)
parser.add_argument('--num_workers', type=int, default=0)
parser.add_argument('--cuda_device', type=int, default=0)

args = vars(parser.parse_args())
args['model'] = 'nco'
# Pretty print the run args
pp.pprint(args)
# hack
args['n_nodes'] = args['input_size']
# Set the random seed
torch.manual_seed(int(args['random_seed']))

torch.cuda.set_device(args['cuda_device'])

# Optionally configure tensorboard
args['run_name'] = args['_id'][-6:] + '-' + args['run_name']
if not args['disable_tensorboard']:
    configure(os.path.join(args['base_dir'], args['log_dir'], args['task'], args['run_name']))

args['test_size'] = args['val_size']
# Task specific configuration - generate dataset if needed
task = args['task'].split('_')
args['COP'] = task[0]  # the combinatorial optimization problem
args, env, training_dataloader, test_dataloader = dataset.build(args, args['epoch_start'])
if args['COP'] == 'mwm2D':
    mwm2D_opt = test_dataloader.dataset.get_average_optimal_weight()

# Open files for writing results
if args['save_stats']:
    fglab_results_dir = os.path.join(args['base_dir'], 'results', 'fglab', args['model'], args['COP'], args['_id'])
    raw_results_dir = os.path.join(args['base_dir'], 'results', 'raw', args['model'], args['COP'], args['_id'])
    try:
        os.makedirs(fglab_results_dir)
        os.makedirs(raw_results_dir)
    except:
        pass
    fglab_results = open(os.path.join(fglab_results_dir, 'scores.json'), 'w')
    raw_results = h5py.File(os.path.join(raw_results_dir, 'raw.hdf5'), 'w')

# Load the model parameters from a saved state
if os.path.exists(args['load_path']):
    print('  [*] Loading model from {}'.format(args['load_path']))

    model = torch.load(
        os.path.join(
            os.getcwd(),
            args['load_path']
        ))
    #model.actor_net.decoder.max_length = args['input_size']
    model.is_train = args['is_train']
else:
    if args['COP'] == 'mwm2D':
        if args['use_decoder']:
            model = MatchingNeuralCombOptRL(
                args['input_size'],
                args['n_features'],
                int(args['embedding_dim']),
                int(args['hidden_dim']),
                args['input_size'], # decoder len
                args['terminating_symbol'],
                int(args['n_glimpses']),
                int(args['n_process_blocks']),
                float(args['tanh_exploration']),
                args['use_tanh'],
                int(args['beam_size']),
                args['is_train'],
                args['use_cuda'])
        else:
            model = MatchingNoDecoder(
                    args['input_size'],
                    args['n_features'],
                    args['embedding_dim'],
                    args['hidden_dim'],
                    args['use_cuda'])
            model.mask_logits = True
    else:
        # Instantiate the Neural Combinatorial Opt with RL module
        model = NeuralCombOptRL(
            args['n_features'],
            int(args['embedding_dim']),
            int(args['hidden_dim']),
            args['input_size'], # decoder len
            args['terminating_symbol'],
            int(args['n_glimpses']),
            int(args['n_process_blocks']),
            float(args['tanh_exploration']),
            args['use_tanh'],
            int(args['beam_size']),
            args['is_train'],
            args['use_cuda'])

args['save_dir'] = os.path.join(args['base_dir'], 'results', 'models', args['model'], args['COP'], args['_id'])
try:
    os.makedirs(args['save_dir'])
except:
    pass

#critic_mse = torch.nn.MSELoss()
#critic_optim = optim.Adam(model.critic_net.parameters(), lr=float(args['critic_net_lr']))
actor_optim = optim.Adam(model.parameters(), lr=float(args['actor_net_lr']))
actor_scheduler = lr_scheduler.MultiStepLR(actor_optim,
        range(int(args['actor_lr_decay_step']), int(args['actor_lr_decay_step']) * 1000,
            int(args['actor_lr_decay_step'])), gamma=float(args['actor_lr_decay_rate']))

#critic_scheduler = lr_scheduler.MultiStepLR(critic_optim,
#        range(int(args['critic_lr_decay_step']), int(args['critic_lr_decay_step']) * 1000,
#            int(args['critic_lr_decay_step'])), gamma=float(args['critic_lr_decay_rate']))

critic_exp_mvg_avg = torch.zeros(1)
beta = args['critic_beta']

if args['use_cuda']:
    model = model.cuda()
    #critic_mse = critic_mse.cuda()
    critic_exp_mvg_avg = critic_exp_mvg_avg.cuda()
step = 0
val_step = 0
tot_R = []
scores = {'_scores': {}}
if not args['is_train']:
    args['n_epochs'] = '1'
epoch = int(args['epoch_start'])

def eval(val_step, final=False):
    # Use (greedy) beam search decoding for validation
    #model.actor_net.decoder.decode_type = "greedy"
    model.decode_type("greedy")
    print('\nstarting eval\n')
    example_input = []
    example_output = []
    eval_R = []
    ratios = []
    optimal = []
    # put in test mode!
    model.eval()
    for batch_id, obs in enumerate(tqdm(test_dataloader,
            disable=args['disable_progress_bar'])):
        obs = Variable(obs, requires_grad=False)
        if args['use_cuda']:
            obs = obs.cuda()
        obs = torch.transpose(obs, 2, 1)
        probs, actions, action_idxs, _ = model(obs)
        if args['COP'] == 'sort':
            R = env(actions, args['use_KT'], args['use_cuda'])
        elif args['COP'] == 'mwm2D':
            # actions is list of len N of (batch_size, n_features)
            x1 = obs[:, :, 0:args['input_size']]
            x2 = torch.stack(actions, 2)
            a = torch.cat([x1, x2], dim=2)
            R = env(torch.transpose(a, 2, 1), args['use_cuda'])
        else:
            R = env(actions, args['use_cuda'])
        eval_R.append(R.data.cpu().numpy())
        val_step += 1
        if val_step % int(args['log_step']) == 0:
            # example_output = []
            # example_input = []
            # for idx, action in enumerate(actions):
            #     if task[0] == 'tsp':
            #         example_output.append(action_idxs[idx][0].data[0])
            #     else:
            #         example_output.append(action[0].data[0])
            #     example_input.append(bat[0, :, idx].data[0])
            #print('Example test input: {}'.format(example_input))
            #print('Example test output: {}'.format(example_output))
            print('step: {}, example reward: {}'.format(val_step, R[0].data[0]))
            # if args['plot_attention']:
            #     probs = torch.cat(probs, 0)
            #     plot_attention(example_input,
            #             example_output, probs.data.cpu().numpy())
        if args['COP'] == 'mwm2D':
            ratios.append(R.data.cpu().numpy() / mwm2D_opt)
    eval_R = np.array(eval_R).ravel()
    #if args['COP'] == 'sort':
    #    # Count how many 1's in eval_R
    #    per_incorrectly_sorted = (len(eval_R) - sum(eval_R == -1.))/(len(eval_R)) * 100.
    #    print('percent incorrectly sorted: {}'.format(per_incorrectly_sorted))
    #    scores['_scores']['percent_incorrectly_sorted_{}'.format(step * args['batch_size'])] = float(per_incorrectly_sorted)
    mean_eval_R = np.mean(eval_R)
    std_eval_R = np.std(eval_R)
    print('Validation overall avg_reward: {}'.format(mean_eval_R))
    print('Validation overall reward std: {}'.format(std_eval_R))
    if not args['disable_tensorboard']:
        log_value('eval_avg_reward', mean_eval_R, val_step)
        log_value('eval_std_reward', std_eval_R, val_step)
    scores['_scores']['eval_avg_reward_{}'.format(step * args['batch_size'])] = float(mean_eval_R)
    #scores['_scores']['eval_std_reward_{}'.format(step * args['batch_size'])] = float(std_eval_R)
    if args['COP'] == 'mwm2D':
        print('Average optimal MWM: {}'.format(mwm2D_opt))
        print('Average optimality ratio: {}'.format(np.mean(ratios)))
        scores['_scores']['optimality_ratio_{}'.format(step * args['batch_size'])] = float(np.mean(ratios))
    model.decode_type("stochastic")
    return val_step

for i in range(epoch, epoch + int(args['n_epochs'])):
    if args['is_train']:
        # eval at 0
        val_step = eval(val_step)
        # put in train mode!
        model.train()
        # sample_batch is [batch_size x input_dim x sourceL]
        for batch_id, obs in enumerate(tqdm(training_dataloader,
                disable=args['disable_progress_bar'])):
            obs = Variable(obs, requires_grad=False)
            if args['use_cuda']:
                obs = obs.cuda()
            obs = torch.transpose(obs, 2, 1)
            probs, actions, actions_idxs, _ = model(obs)
            if args['COP'] == 'sort':
                R = env(actions, args['use_KT'], args['use_cuda'])
            elif args['COP'] == 'mwm2D':
                # actions is list of len N of (batch_size, n_features)
                x1 = obs[:, :, 0:args['input_size']]
                x2 = torch.stack(actions, 2)
                a = torch.cat([x1, x2], dim=2)
                R = env(torch.transpose(a, 2, 1), args['use_cuda'])
            else:
                R = env(actions, args['use_cuda'])
            tot_R.append(R.data.cpu().numpy())
            if batch_id == 0:
                critic_exp_mvg_avg = R.mean()
            else:
                critic_exp_mvg_avg = (critic_exp_mvg_avg * beta) + ((1. - beta) * R.mean())
            advantage = R - critic_exp_mvg_avg
            if not args['use_decoder']:
                logprobs = torch.stack(probs).sum(dim=0)
                nll = -logprobs.detach()
            else:
                logprobs = 0
                nll = 0
                for prob in probs:
                    # compute the sum of the log probs
                    # for each tour in the batch
                    logprob = torch.log(prob)
                    nll += -logprob.detach()
                    logprobs = logprobs + logprob
                # guard against nan
                #nll[nll != nll] = 0.
                # clamp any -inf's to 0 to throw away this tour
                #logprobs[logprobs < -1000] = 0.
                # multiply each time step by the advanrate
            reinforce = advantage.detach() * logprobs
            actor_loss = reinforce.mean()
            actor_optim.zero_grad()
            actor_loss.backward()
            # clip gradient norms
            torch.nn.utils.clip_grad_norm(model.parameters(),
                    float(args['max_grad_norm']), norm_type=2)
            actor_optim.step()
            actor_scheduler.step()
            critic_exp_mvg_avg = critic_exp_mvg_avg.detach()
            #critic_scheduler.step()
            #R = R.detach()
            #critic_loss = critic_mse(v.squeeze(1), R)
            #critic_optim.zero_grad()
            #critic_loss.backward()
            #torch.nn.utils.clip_grad_norm(model.critic_net.parameters(),
            #        float(args['max_grad_norm']), norm_type=2)
            #critic_optim.step()
            step += 1
            if not args['disable_tensorboard']:
                log_value('Running_avg_reward', -1 * R.mean().data[0], step)
                log_value('actor_loss', actor_loss.data[0], step)
                #log_value('critic_loss', critic_loss.data[0], step)
                log_value('critic_exp_mvg_avg', critic_exp_mvg_avg.data[0], step)
                log_value('nll', nll.mean().data[0], step)
            if step % int(args['log_step']) == 0:
                print('epoch: {}, train_batch_id: {}, avg_reward: {}'.format(
                    i, batch_id, R.mean().data[0]))
                example_output = []
                #example_input = []
                if args['COP'] == 'sort':
                    for idx, action in enumerate(actions):
                        example_output.append(round(action[0].data[0]))  # <-- ??
                    #if task[0] == 'tsp':
                    #    example_output.append(actions_idxs[idx][0].data[0])
                    #else:
                    #example_input.append(sample_batch[0, :, idx][0])
                    #print('Example train input: {}'.format(example_input))
                    print('Example train output: {}'.format(example_output))
        if args['save_model']:
            print(' [*] saving model...')
            torch.save(model, os.path.join(args['save_dir'], 'nco-COP-{}-N-{}-epoch-{}.pt'.format(args['COP'], args['input_size'], i)))
# Eval one last time
val_step = eval(val_step, True)

if args['save_stats']:
    # write training stats to file
    json.dump(scores, fglab_results)
    tot_R = np.array(tot_R).ravel()
    raw_results.create_dataset('training_rewards', data=tot_R)
    # close files
    fglab_results.close()
    raw_results.close()


import numpy as np
from scipy import interpolate
import scipy
import astropy
import matplotlib
from datetime import datetime
from scipy.stats.mstats import mquantiles
import os
from astropy.io import fits
import glob
from parameters import *
import sys
import argparse
from astropy.table import Table
#import datetime
##sys.path.insert(0, 'python')
# matplotlib.rc('text', usetex=True)
# matplotlib.rcParams['text.latex.preamble'] = [r"\usepackage{amsmath}",
#                                    r"\usepackage{color}"]
# matplotlib.use('Agg')
from ptemcee import Sampler as PTSampler
import matplotlib.pyplot as plt
import matplotlib
##import corner
startTime = datetime.now()

parser = argparse.ArgumentParser(description="Calculate radial velocities and stellar parameters of WEAVE target spectra.")

parser.add_argument("--infiles", type=str, required=True, help="input file", nargs=1)
parser.add_argument("--outdir", type=str, required=True, help="output directory", nargs=1)
parser.add_argument("--targlist", type=str, required=True, help="path to list of FIBREIDs or TARGIDs to be analysed. To analyse all BA stars, enter 'all'.", nargs=1)
parser.add_argument("--params", type=str, required=False, default='parameters.py', help="path to parameter file", nargs=1)
parser.add_argument("--apsclass", type=bool, required=False, default=False, help="to read APS classification", nargs=1)
parser.add_argument("--override", type=bool, required=False, default=False, help="remove existing files and start from scratch", nargs=1)
#parser.add_argument("--setups", type=str, default=None, required=True, help="input setups", nargs='*')

args = parser.parse_args()
write_directory = args.outdir[0]
target_list = args.targlist[0]
data_file = args.infiles[0]
apsclassification=args.apsclass#[0]
override=args.override#[0]

if apsclassification==True:
	print('WARNING: Code still not ready for reading L2 and check APS classification')
	print('		Please, set --apsclass False')
	sys.exit()

# processing (cropping, smoothing, rebinning, rotational broadening) templates if required

if process_templates==True:
	print('Beginning processing of templates')
	from PyAstronomy import pyasl
	import scipy.stats
	templatelist = np.genfromtxt(template_list_file, dtype=None, encoding=None)

	# Ensuring the number of templates to be processed matches the number of points on the grid.
	notemps = len(templatelist)
	d_size = len(Teff) * len(logg) * len(vsini)
	if notemps < d_size:
		print('ERROR: The amount of templates to be processed is less than the amount of points in the grid. Add more templates to the list or reduce the grid as necessary (see PARAMETER BOUNDARIES in parameters file).')
		sys.exit()
	if notemps > d_size:
		print('ERROR: The amount of templates to be processed is greater than the amount of points in the grid. Remove templates from the list or expand the grid as necessary (see PARAMETER BOUNDARIES in parameters file).')
		sys.exit()

	def restrict_range(w, f):		# to crop the templates
		our_range = (w > template_crop_min_wav) & (w < template_crop_max_wav)
		w, f = w[our_range], f[our_range]
		return w, f

	def smooth(w, f, sig):		# to smooth/broaden templates to match resolution of observed spectrum
		f = pyasl.broadGaussFast(w, f, sig)
		return f

	def rotbroad(w, f, vsini):		# to rotationally broaden the templates
		f = pyasl.rotBroad(w, f, 0.6, float(vsini))
		return f

	def rebin(w, f, samp):		# to rebin the templates to match the sampling of the observed spectrum
		f, bin_edges, binnumber = scipy.stats.binned_statistic(w, f, statistic = 'mean', bins = (w[-1] - w[0]) / samp)
		bin_width = bin_edges[1] - bin_edges[0]
		w = bin_edges[1:] - bin_width/2
		return w, f

	def process_template(w, f, sig=sigma, samp=sampling):
		if sigma != None:
			f = smooth(w, f, sig)
		if sampling != None:
			w, f = rebin(w, f, samp)
		file = open(template_write_directory + '/'+ os.path.splitext(os.path.basename(i))[0] + '_processed', "w")
		for index in range(len(w)):
			file.write(str(w[index]) + " " + str(f[index]) + "\n")
		file.close()
		return w, f

	def process_template_vsini(w, f, vsini, sig=sigma, samp=sampling):
		f = rotbroad(w, f, vsini)
		if sigma != None:
			f = smooth(w, f, sig)
		if sampling != None:
			w, f = rebin(w, f, samp)
		file = open(template_write_directory + '/' + os.path.splitext(os.path.basename(i))[0] + '_processed_vsini' + str(int(vsini)), "w")
		for index in range(len(w)):
			file.write(str(w[index]) + " " + str(f[index]) + "\n")
		file.close()
		return w, f

	def write_to_grid(templatename, f, t_ind=temp_ind, l_ind=logg_ind, v_ind=vsini_ind, vsini0=None):
		# finding the teff and logg information from the template name, to write to corresponding grid point
		t1 = float(templatename[temp_ind[0]:temp_ind[1]])
		l1 = float(templatename[logg_ind[0]:logg_ind[1]])
		if vsini0==None:
			v1 = float(templatename[vsini_ind[0]:vsini_ind[1]])
		else:
			v1 = float(vsini0)

		print('Teff: ' + str(t1) + ', logg: ' + str(l1) + ', vsini: ' + str(v1))

		d[np.where(Teff == t1)[0][0], np.where(logg == l1)[0][0], np.where(vsini == v1)[0][0]] = f

	t_wavelength_0, t_flux_0 = np.genfromtxt(templatelist[0], unpack=True, usecols=(wav_ind, flux_ind), dtype=None, encoding=None)
	t_wavelength_0, t_flux_0 = restrict_range(t_wavelength_0, t_flux_0)
	if sampling != None:
		t_wavelength_0, t_flux_0 = rebin(t_wavelength_0, t_flux_0, sampling)

	d = np.zeros((len(Teff), len(logg), len(vsini), len(t_wavelength_0)))
	d_filled = 0

	for i in templatelist:
		print('Processing: ' + i)
		t_wavelength, t_flux = np.genfromtxt(i, unpack=True, usecols=(wav_ind, flux_ind), dtype=None, encoding=None)
		t_wavelength, t_flux = restrict_range(t_wavelength, t_flux)
		if rotbroads == None:
			_, t_flux = process_template(t_wavelength, t_flux, sigma, sampling)
			print('template processed, writing to grid at:')
			write_to_grid(i, t_flux)
			d_filled += 1
			print('grid is ' + str("{0:.1f}".format(float(d_filled)/float(d_size) * 100.)) + ' percent full')

		else:
			print('processing for vsini: 0')
			_, t_flux_v0 = process_template(t_wavelength, t_flux, sig=sigma, samp=sampling)
			print('template processed, writing to grid at:')
			write_to_grid(i, t_flux_v0, vsini0=0.)
			d_filled += 1
			print('grid is ' + str("{0:.1f}".format(float(d_filled)/float(d_size) * 100.)) + ' percent full')

			for j in rotbroads:
				if j == 0:
					continue
				print('processing for vsini: ' + str(j))
				_, t_flux_vj = process_template_vsini(t_wavelength, t_flux, j, sig=sigma, samp=sampling)
				print('template processed, writing to grid at:')
				write_to_grid(i, t_flux_vj, vsini0=j)
				d_filled += 1
				print('grid is ' + str("{0:.1f}".format(float(d_filled)/float(d_size) * 100.)) + ' percent full')

	np.save(template_write_directory + 'template_grid.npy', d)

# creating /plots and /results folders in write_directory if they dont already exist
if not os.path.exists(write_directory + '/plots'):
    os.mkdir(write_directory + '/plots')
    print(write_directory + 'plots folder created')
if not os.path.exists(write_directory + '/results'):
    os.mkdir(write_directory + '/results')
    print(write_directory + 'results folder created')

if process_templates==False:
	# location of folder containing this script
	templatedirectory = os.path.dirname(os.path.abspath(__file__)) + '/templates/'
	# loading the template flux grid
	flux_data_all = np.load(templatedirectory + 'template_grid.npy')
	#loading the wavelength data
	templatewavelength = np.loadtxt(templatedirectory + 'wavelength_data.dat')
else:
	# the first template for the template wavelength data
	flux_data_all = np.load(template_write_directory + 'template_grid.npy')
	templatewavelength = t_wavelength_0

template_min_wav = min_wav*(1.0 + (min(drv)/299792.458)) - 10.
template_max_wav = max_wav*(1.0 + (max(drv)/299792.458)) + 10.

templatemask = (templatewavelength > template_min_wav) & (templatewavelength < template_max_wav)
flux_data = np.zeros((len(Teff), len(logg), len(vsini), len(templatewavelength[templatemask])))
templatewavelength = templatewavelength[templatemask]

for ii in range(len(Teff)):
	for jj in range(len(logg)):
		for kk in range(len(vsini)):
			flux_data[ii,jj,kk] = flux_data_all[ii,jj,kk][templatemask]

ipo = interpolate.RegularGridInterpolator((Teff, logg, vsini), flux_data, method='linear')

# grid of parameters for walker initial positions
ini_grid = [Teff, logg, vsini, drv, slopes, intercepts]
ndim=6

# define edges of parameter space
teffmin, teffmax, loggmin, loggmax, vsinimin, vsinimax, rvmin, rvmax, slopemin, slopemax, interceptmin, interceptmax = min(Teff), max(Teff), min(logg), max(logg), min(vsini), max(vsini), min(drv), max(drv), min(slopes), max(slopes), min(intercepts), max(intercepts)

def model(X, wavelength):
	i, j, k, l, m, n = X
	# interpolating template grid with teff (i), logg (j), vsini (k) trial parameter
	templateflux = ipo([i, j, k])[0]
	# interpolating on wavelength axis for trial RV (l)
	fi = interpolate.interp1d(templatewavelength*(1.0 + l/299792.458), templateflux)
	return fi(wavelength)

def lnprior(X):
	i, j, k, l, m, n = X
	# flat prior, edges should corespond to template grid
	if (teffmin <= i <= teffmax) & (loggmin <= j <= loggmax) & (vsinimin <= k <= vsinimax) & (rvmin <= l <= rvmax) & (slopemin <= m <= slopemax) & (interceptmin <= n <= interceptmax):
		return 0.0
	else:
		return -np.inf

def lnlike(X, wavelength, flux, noisespec, mask):
	i, j, k, l, m, n = X
	z = m, n
	f = np.poly1d(z)
	if exclude_region == False:
		return -(np.sum((flux - (model(X, wavelength)*f(wavelength)))**2/(2*PPRE*noisespec**2)))
	else:
		return -(np.sum((flux[mask] - (model(X, wavelength)*f(wavelength))[mask])**2/(2*PPRE*noisespec[mask]**2)))

def mcmc_one(t):

	print("Processing: " + t)
	targ_start = datetime.now()

	# checking if spectrum file exists, and if result already written
	if os.path.exists(write_directory + '/results/' + os.path.splitext(os.path.basename(t))[0] + '_results'):
		if override:
			os.remove(write_directory + '/results/' + os.path.splitext(os.path.basename(t))[0] + '_results')
			os.remove(write_directory + '/results/' + os.path.splitext(os.path.basename(t))[0] + '_spec')
		else:
			print('WARNING: result for '+t+' already exists, skipping')
			return

	# specify unnormalised calibrated flux
	with fits.open(data_file) as ALLDATA:
		final_spectra = ALLDATA['RED_DATA'].data[info['TARGID'] == t][0]
		spectra_before_sky_subtraction = ALLDATA['RED_DATA_NOSS'].data[info['TARGID'] == t][0]
		Calibration_function = ALLDATA['RED_SENSFUNC'].data[info['TARGID'] == t][0]
		try:
			idx=list(ALLDATA['FIBTABLE'].data['FIBREID']).index(t)
		except:
			idx=list(ALLDATA['FIBTABLE'].data['TARGID']).index(t)
		NSPEC = ALLDATA['FIBTABLE'].data['Nspec'][idx]
		FIBREID = ALLDATA['FIBTABLE'].data['FIBREID'][idx]
		CNAME = ALLDATA['FIBTABLE'].data['CNAME'][idx]
	flux = final_spectra * Calibration_function * 1.0e18

	# restrict to desired wavelength range
	flux = flux[targetmask]

	# create corresponsing noise spectrum
	noisespec = np.sqrt((2.*spectra_before_sky_subtraction - final_spectra)*Calibration_function*1.0e18)
	noisespec = noisespec[targetmask]

	# mask for excluding a wavelength region
	if exclude_region == True:
		mask = np.zeros(len(wavelength), dtype=bool)
		for line in lines:
			mask |= (wavelength >= line[0]) & (wavelength <= line[1])
	else:
		mask = np.ones(len(wavelength), dtype=bool)

	# choose initial walker positions
	pos = [[[np.random.choice(i) for i in ini_grid] for i in range(nwalkers)] for i in range(ntemps)]

	# initialise MCMC sampler
	sampler = PTSampler(ntemps=ntemps, nwalkers=nwalkers, dim=ndim, logl=lnlike, logp=lnprior, Tmax=np.inf, loglargs=[wavelength, flux, noisespec, mask])

	# run MCMC sampler for burn period
	if progress_bar==True:
		print("running burn")
	pos, prob, state = sampler.run_mcmc(pos, burn, adapt=True, progress_bar=progress_bar)
	# reset sampler, run MCMC sampler for run period with walkers starting at their positions at the end of burn
	sampler.reset()
	if progress_bar==True:
		print("running runs")
	sampler.run_mcmc(pos, runs, adapt=True, progress_bar=progress_bar)
	samples=sampler.chain[0, :, :, :].reshape((-1, ndim))

#	# plot walker paths
#	ylabels = ['Teff', 'logg', 'vsini', 'RV', 'slope', 'intercept']
#	for m in range(ndim):
#		plt.subplot(ndim,1,m+1)
#		plt.plot(sampler.chain[0,:,:,m].transpose(), alpha=0.2)
#		plt.ylabel(ylabels[m])
#	plt.xlabel('Step')
#	plt.savefig(write_directory + '/plots/' + t + '_walkers.png', bbox_inches='tight')
#	plt.close()

	# calculate 16th, 50th, 84th quantiles of the parameter samples
	quantiles = mquantiles(samples, prob=[0.16, 0.50, 0.84], axis=0)

	targetname = os.path.splitext(os.path.basename(t))[0]
	acceptance_r = np.mean(sampler.acceptance_fraction)

	# print acceptance fraction, should be between 0.2-0.5 for efficient sampling
	print("Mean acceptance fraction: {0:.3f}"
                .format(acceptance_r))

	# The parameter results
	Teff_r = quantiles[1][0]
	Teffminus_r = quantiles[1][0] - quantiles[0][0]
	Teffplus_r = quantiles[2][0] - quantiles[1][0]
	logg_r = quantiles[1][1]
	loggminus_r = quantiles[1][1] - quantiles[0][1]
	loggplus_r = quantiles[2][1] - quantiles[1][1]
	vsini_r = quantiles[1][2]
	vsiniminus_r = quantiles[1][2] - quantiles[0][2]
	vsiniplus_r = quantiles[2][2] - quantiles[1][2]
	RV_r = quantiles[1][3]
	RVminus_r = quantiles[1][3] - quantiles[0][3]
	RVplus_r = quantiles[2][3] - quantiles[1][3]
	slope_r = quantiles[1][4]
	slopeminus_r = quantiles[1][4] - quantiles[0][4]
	slopeplus_r = quantiles[2][4] - quantiles[1][4]
	intercept_r = quantiles[1][5]
	interceptminus_r = quantiles[1][5] - quantiles[0][5]
	interceptplus_r = quantiles[2][5] - quantiles[1][5]

	#fig = corner.corner(samples, quantiles=[0.16, 0.50, 0.84], labels=['Teff', 'log(g)', 'vsini', 'RV', 'slope', 'intercept'], show_titles=True, title_kwargs={"fontsize": 10}, plot_datapoints=True, plot_contours=True, auto_bars=True, data_kwargs={"alpha": 0.005})
	# fig.savefig(write_directory + t + '_cornerplot.png', bbox_inches='tight')
	# plt.close()

	# parameters of best fit (for plotting)
	Xp = Teff_r, logg_r, vsini_r, RV_r, slope_r, intercept_r
	fp = np.poly1d(Xp[-2:])
	fitp = fp(wavelength)

#	# plot spectrum and best-fit
#	plt.plot(wavelength, flux, wavelength, model(Xp, wavelength) * fitp)
#	if exclude_region == True:
#		for n,line in enumerate(lines):
#			if n == 0:
#				plt.vlines([line[1]], flux.min()*0.8, flux.max()*1.2, colors='r', linestyles='dashed')
#			elif n == (len(lines)-1):
#				plt.vlines([line[0]], flux.min()*0.8, flux.max()*1.2, colors='r', linestyles='dashed')
#			else:
#				plt.vlines([line[0], line[1]], flux.min()*0.8, flux.max()*1.2, colors='r', linestyles='dashed')

#	plt.xlim(min_wav, max_wav)
#	plt.ylim(flux.min()*0.8, flux.max()*1.2)
#	plt.xlabel(r'Wavelength ($\AA$)')
#	plt.ylabel('Calibrated counts')
#	plt.savefig(write_directory + '/plots/' + t + '_spectrum.png', bbox_inches='tight')
#	plt.close()

#	# plot mapping function
#	plt.plot(wavelength, fitp, wavelength, flux / model(Xp, wavelength))
#	plt.xlim(min_wav, max_wav)
#	plt.xlabel(r'Wavelength ($\AA$)')
#	plt.ylabel('Calibrated counts')
#	plt.savefig(write_directory + '/plots/' + t + '_mapping_function.png', bbox_inches='tight')
#	plt.close()

	# write results
	tab = open(write_directory + '/results/' + t +  '_results', "w")
	tab.write(np.str(NSPEC) + " " + np.str(FIBREID) + " " + np.str(CNAME) + " " + t + " " + np.str(acceptance_r) + " " + np.str(Teff_r) + " " + np.str(Teffminus_r) + " " + np.str(Teffplus_r) + " " + np.str(logg_r) + " " + np.str(loggminus_r) + " " + np.str(loggplus_r) + " " + np.str(vsini_r) + " " + np.str(vsiniminus_r) + " " + np.str(vsiniplus_r) + " " + np.str(RV_r) + " " + np.str(RVminus_r) + " " + np.str(RVplus_r) + " " + np.str(slope_r) + " " + np.str(slopeminus_r) + " " + np.str(slopeplus_r) + " " + np.str(intercept_r) + " " + np.str(interceptminus_r) + " " + np.str(interceptplus_r) + "\n")
	tab.close()

	tab = open(write_directory + '/results/' + t +  '_spec', "w")
	tab.write("{0} ; {1} ; {2}; {3} ; {4} ; {5}\n".format(NSPEC,FIBREID,CNAME,t,list(flux),list(model(Xp, wavelength) * fitp)))
	tab.close()

	targ_end = datetime.now() - targ_start
	print(targ_end)

def modheader(hdul):
	hdul[1].header.comments['TTYPE1']='The number of the spectrum'
	hdul[1].header.comments['TTYPE2']='Fibre id'
	hdul[1].header.comments['TTYPE3']='WEAVE object name from coordinates'
	hdul[1].header.comments['TTYPE4']='Identifier of the target assigned by survey'
	hdul[1].header.comments['TTYPE5']='Mean fraction of proposed walker jumps'
	hdul[1].header.comments['TTYPE6']='Effective temperature'
	hdul[1].header.comments['TTYPE7']='1-sigma negative uncertainty on Teff'
	hdul[1].header.comments['TTYPE8']='1-sigma positive uncertainty on Teff'
	hdul[1].header.comments['TTYPE9']='Surface gravity log(g)'
	hdul[1].header.comments['TTYPE10']='1-sigma negative uncertainty on logg'
	hdul[1].header.comments['TTYPE11']='1-sigma positive uncertainty on logg'
	hdul[1].header.comments['TTYPE12']='Projected rotational velocity'
	hdul[1].header.comments['TTYPE13']='1-sigma negative uncertainty on vsini'
	hdul[1].header.comments['TTYPE14']='1-sigma positive uncertainty on vsini'
	hdul[1].header.comments['TTYPE15']='Radial/line-of-sight velocity'
	hdul[1].header.comments['TTYPE16']='1-sigma negative uncertainty on RV'
	hdul[1].header.comments['TTYPE17']='1-sigma positive uncertainty on RV'
	hdul[1].header.comments['TTYPE18']='Slope of mapping function '
	hdul[1].header.comments['TTYPE19']='1-sigma negative uncertainty on slope'
	hdul[1].header.comments['TTYPE20']='1-sigma positive uncertainty on slope'
	hdul[1].header.comments['TTYPE21']='Intercept of mapping function'
	hdul[1].header.comments['TTYPE22']='1-sigma negative uncertainty on intercept'
	hdul[1].header.comments['TTYPE23']='1-sigma positive uncertainty on intercept'

	hdul[1].header.comments['TFORM1']='data format of field: integer'
	hdul[1].header.comments['TFORM2']='data format of field: integer'
	hdul[1].header.comments['TFORM3']='data format of field: ASCII Character'
	hdul[1].header.comments['TFORM4']='data format of field: ASCII Character'
	hdul[1].header.comments['TFORM5']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM6']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM7']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM8']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM9']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM10']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM11']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM12']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM13']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM14']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM15']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM16']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM17']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM18']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM19']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM20']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM21']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM22']='data format of field: 4-byte REAL'
	hdul[1].header.comments['TFORM23']='data format of field: 4-byte REAL'

	hdul[1].header.set('TDISP1','I4',after='TFORM1')
	hdul[1].header.set('TDISP2','I4',after='TFORM2')
	hdul[1].header.set('TDISP3','A20',after='TFORM3')
	hdul[1].header.set('TDISP4','A30',after='TFORM4')
	hdul[1].header.set('TDISP5','F7.3',after='TFORM5')
	hdul[1].header.set('TDISP6','F7.3',after='TFORM6')
	hdul[1].header.set('TDISP7','F7.3',after='TFORM7')
	hdul[1].header.set('TDISP8','F7.3',after='TFORM8')
	hdul[1].header.set('TDISP9','F7.3',after='TFORM9')
	hdul[1].header.set('TDISP10','F7.3',after='TFORM10')
	hdul[1].header.set('TDISP11','F7.3',after='TFORM11')
	hdul[1].header.set('TDISP12','F7.3',after='TFORM12')
	hdul[1].header.set('TDISP13','F7.3',after='TFORM13')
	hdul[1].header.set('TDISP14','F7.3',after='TFORM14')
	hdul[1].header.set('TDISP15','F7.3',after='TFORM15')
	hdul[1].header.set('TDISP16','F7.3',after='TFORM16')
	hdul[1].header.set('TDISP17','F7.3',after='TFORM17')
	hdul[1].header.set('TDISP18','F7.3',after='TFORM18')
	hdul[1].header.set('TDISP19','F7.3',after='TFORM19')
	hdul[1].header.set('TDISP20','F7.3',after='TFORM20')
	hdul[1].header.set('TDISP21','F7.3',after='TFORM21')
	hdul[1].header.set('TDISP22','F7.3',after='TFORM22')
	hdul[1].header.set('TDISP23','F7.3',after='TFORM23')

	hdul[1].header.set('TUCD1','meta.id',after='TDISP1')
	hdul[1].header.set('TUCD2','meta.id',after='TDISP2')
	hdul[1].header.set('TUCD3','meta.id;meta.main',after='TDISP3')
	hdul[1].header.set('TUCD4','meta.id',after='TDISP4')
	hdul[1].header.set('TUCD5','obs.param',after='TDISP5')
	hdul[1].header.set('TUCD6','phys.temperature.effective',after='TDISP6')
	hdul[1].header.set('TUCD7','stat.error;phys.temperature.effective',after='TDISP7')
	hdul[1].header.set('TUCD8','stat.error;phys.temperature.effective',after='TDISP8')
	hdul[1].header.set('TUCD9','phys.gravity',after='TDISP9')
	hdul[1].header.set('TUCD10','stat.error;phys.gravity',after='TDISP10')
	hdul[1].header.set('TUCD11','stat.error;phys.gravity',after='TDISP11')
	hdul[1].header.set('TUCD12','phys.veloc.rotat',after='TDISP12')
	hdul[1].header.set('TUCD13','stat.error;phys.veloc.rotat',after='TDISP13')
	hdul[1].header.set('TUCD14','stat.error;phys.veloc.rotat',after='TDISP14')
	hdul[1].header.set('TUCD15','spect.dopplerVeloc',after='TDISP15')
	hdul[1].header.set('TUCD16','stat.error;spect.dopplerVeloc',after='TDISP16')
	hdul[1].header.set('TUCD17','stat.error;spect.dopplerVeloc',after='TDISP17')
	hdul[1].header.set('TUCD18','obs.param',after='TDISP18')
	hdul[1].header.set('TUCD19','stat.error;obs.param',after='TDISP19')
	hdul[1].header.set('TUCD20','stat.error;obs.param',after='TDISP20')
	hdul[1].header.set('TUCD21','obs.param',after='TDISP21')
	hdul[1].header.set('TUCD22','stat.error;obs.param',after='TDISP22')
	hdul[1].header.set('TUCD23','stat.error;obs.param',after='TDISP23')

	hdul[1].header.set('TUNIT6','K',after='TUCD6')
	hdul[1].header.set('TUNIT7','K',after='TUCD7')
	hdul[1].header.set('TUNIT8','K',after='TUCD8')
	hdul[1].header.set('TUNIT12','km/s',after='TUCD12')
	hdul[1].header.set('TUNIT13','km/s',after='TUCD13')
	hdul[1].header.set('TUNIT14','km/s',after='TUCD14')
	hdul[1].header.set('TUNIT15','km/s',after='TUCD15')
	hdul[1].header.set('TUNIT16','km/s',after='TUCD16')
	hdul[1].header.set('TUNIT17','km/s',after='TUCD17')

	hdul[1].header.comments['TDISP1']='Display format for column'
	hdul[1].header.comments['TDISP2']='Display format for column'
	hdul[1].header.comments['TDISP3']='Display format for column'
	hdul[1].header.comments['TDISP4']='Display format for column'
	hdul[1].header.comments['TDISP5']='Display format for column'
	hdul[1].header.comments['TDISP6']='Display format for column'
	hdul[1].header.comments['TDISP7']='Display format for column'
	hdul[1].header.comments['TDISP8']='Display format for column'
	hdul[1].header.comments['TDISP9']='Display format for column'
	hdul[1].header.comments['TDISP10']='Display format for column'
	hdul[1].header.comments['TDISP11']='Display format for column'
	hdul[1].header.comments['TDISP12']='Display format for column'
	hdul[1].header.comments['TDISP13']='Display format for column'
	hdul[1].header.comments['TDISP14']='Display format for column'
	hdul[1].header.comments['TDISP15']='Display format for column'
	hdul[1].header.comments['TDISP16']='Display format for column'
	hdul[1].header.comments['TDISP17']='Display format for column'
	hdul[1].header.comments['TDISP18']='Display format for column'
	hdul[1].header.comments['TDISP19']='Display format for column'
	hdul[1].header.comments['TDISP20']='Display format for column'
	hdul[1].header.comments['TDISP21']='Display format for column'
	hdul[1].header.comments['TDISP22']='Display format for column'
	hdul[1].header.comments['TDISP23']='Display format for column'

	hdul[1].header.comments['TUCD1']='UCD for column'
	hdul[1].header.comments['TUCD2']='UCD for column'
	hdul[1].header.comments['TUCD3']='UCD for column'
	hdul[1].header.comments['TUCD4']='UCD for column'
	hdul[1].header.comments['TUCD5']='UCD for column'
	hdul[1].header.comments['TUCD6']='UCD for column'
	hdul[1].header.comments['TUCD7']='UCD for column'
	hdul[1].header.comments['TUCD8']='UCD for column'
	hdul[1].header.comments['TUCD9']='UCD for column'
	hdul[1].header.comments['TUCD10']='UCD for column'
	hdul[1].header.comments['TUCD11']='UCD for column'
	hdul[1].header.comments['TUCD12']='UCD for column'
	hdul[1].header.comments['TUCD13']='UCD for column'
	hdul[1].header.comments['TUCD14']='UCD for column'
	hdul[1].header.comments['TUCD15']='UCD for column'
	hdul[1].header.comments['TUCD16']='UCD for column'
	hdul[1].header.comments['TUCD17']='UCD for column'
	hdul[1].header.comments['TUCD18']='UCD for column'
	hdul[1].header.comments['TUCD19']='UCD for column'
	hdul[1].header.comments['TUCD20']='UCD for column'
	hdul[1].header.comments['TUCD21']='UCD for column'
	hdul[1].header.comments['TUCD22']='UCD for column'
	hdul[1].header.comments['TUCD23']='UCD for column'

	hdul[1].header.comments['TUNIT6']='physical unit of field'
	hdul[1].header.comments['TUNIT7']='physical unit of field'
	hdul[1].header.comments['TUNIT8']='physical unit of field'
	hdul[1].header.comments['TUNIT12']='physical unit of field'
	hdul[1].header.comments['TUNIT13']='physical unit of field'
	hdul[1].header.comments['TUNIT14']='physical unit of field'
	hdul[1].header.comments['TUNIT15']='physical unit of field'
	hdul[1].header.comments['TUNIT16']='physical unit of field'
	hdul[1].header.comments['TUNIT17']='physical unit of field'

	hdul[1].header.set('TPROP1',0,after='TUCD1')
	hdul[1].header.set('TPROP2',0,after='TUCD2')
	hdul[1].header.set('TPROP3',0,after='TUCD3')
	hdul[1].header.set('TPROP4',0,after='TUCD4')
	hdul[1].header.set('TPROP5',0,after='TUCD5')
	hdul[1].header.set('TPROP6',0,after='TUCD6')
	hdul[1].header.set('TPROP7',0,after='TUCD7')
	hdul[1].header.set('TPROP8',0,after='TUCD8')
	hdul[1].header.set('TPROP9',0,after='TUCD9')
	hdul[1].header.set('TPROP10',0,after='TUCD10')
	hdul[1].header.set('TPROP11',0,after='TUCD11')
	hdul[1].header.set('TPROP12',0,after='TUCD12')
	hdul[1].header.set('TPROP13',0,after='TUCD13')
	hdul[1].header.set('TPROP14',0,after='TUCD14')
	hdul[1].header.set('TPROP15',0,after='TUCD15')
	hdul[1].header.set('TPROP16',0,after='TUCD16')
	hdul[1].header.set('TPROP17',0,after='TUCD17')
	hdul[1].header.set('TPROP18',0,after='TUCD18')
	hdul[1].header.set('TPROP19',0,after='TUCD19')
	hdul[1].header.set('TPROP20',0,after='TUCD20')
	hdul[1].header.set('TPROP21',0,after='TUCD21')
	hdul[1].header.set('TPROP22',0,after='TUCD22')
	hdul[1].header.set('TPROP23',0,after='TUCD23')

	hdul[1].header.comments['TPROP1']='Public column'
	hdul[1].header.comments['TPROP2']='Public column'
	hdul[1].header.comments['TPROP3']='Public column'
	hdul[1].header.comments['TPROP4']='Public column'
	hdul[1].header.comments['TPROP5']='Public column'
	hdul[1].header.comments['TPROP6']='Public column'
	hdul[1].header.comments['TPROP7']='Public column'
	hdul[1].header.comments['TPROP8']='Public column'
	hdul[1].header.comments['TPROP9']='Public column'
	hdul[1].header.comments['TPROP10']='Public column'
	hdul[1].header.comments['TPROP11']='Public column'
	hdul[1].header.comments['TPROP12']='Public column'
	hdul[1].header.comments['TPROP13']='Public column'
	hdul[1].header.comments['TPROP14']='Public column'
	hdul[1].header.comments['TPROP15']='Public column'
	hdul[1].header.comments['TPROP16']='Public column'
	hdul[1].header.comments['TPROP17']='Public column'
	hdul[1].header.comments['TPROP18']='Public column'
	hdul[1].header.comments['TPROP19']='Public column'
	hdul[1].header.comments['TPROP20']='Public column'
	hdul[1].header.comments['TPROP21']='Public column'
	hdul[1].header.comments['TPROP22']='Public column'
	hdul[1].header.comments['TPROP23']='Public column'

	hdul[1].header.set('TDMIN1',1,after='TUCD1')
	hdul[1].header.set('TDMIN2',1,after='TUCD2')
	hdul[1].header.set('TDMIN5',0,after='TUCD5')
	hdul[1].header.set('TDMIN6',0,after='TUCD6')
	hdul[1].header.set('TDMIN7',0,after='TUCD7')
	hdul[1].header.set('TDMIN8',0,after='TUCD8')
	hdul[1].header.set('TDMIN9',0,after='TUCD9')
	hdul[1].header.set('TDMIN10',0,after='TUCD10')
	hdul[1].header.set('TDMIN11',0,after='TUCD11')
	hdul[1].header.set('TDMIN12',0,after='TUCD12')
	hdul[1].header.set('TDMIN13',0,after='TUCD13')
	hdul[1].header.set('TDMIN14',0,after='TUCD14')
	hdul[1].header.set('TDMIN16',0,after='TUCD16')
	hdul[1].header.set('TDMIN17',0,after='TUCD17')
	hdul[1].header.set('TDMIN18',-1,after='TUCD18')
	hdul[1].header.set('TDMIN19',0,after='TUCD19')
	hdul[1].header.set('TDMIN20',0,after='TUCD20')
	hdul[1].header.set('TDMIN22',0,after='TUCD22')
	hdul[1].header.set('TDMIN23',0,after='TUCD23')

	hdul[1].header.comments['TDMIN1']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN2']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN5']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN6']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN7']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN8']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN9']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN10']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN11']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN12']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN13']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN14']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN16']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN17']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN18']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN19']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN20']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN22']='Minimum value expected for field'
	hdul[1].header.comments['TDMIN23']='Minimum value expected for field'

	hdul[1].header.set('TDMAX1',960,after='TDMIN1')
	hdul[1].header.set('TDMAX2',1100,after='TDMIN2')
	hdul[1].header.set('TDMAX5',1,after='TDMIN5')
	hdul[1].header.set('TDMAX18',1,after='TDMIN18')

	hdul[1].header.comments['TDMAX1']='Maximum value expected for field'
	hdul[1].header.comments['TDMAX2']='Maximum value expected for field'
	hdul[1].header.comments['TDMAX5']='Maximum value expected for field'
	hdul[1].header.comments['TDMAX18']='Maximum value expected for field'

	hdul[2].header.comments['TTYPE1']='The number of the spectrum'
	hdul[2].header.comments['TTYPE2']='Fibre id'
	hdul[2].header.comments['TTYPE3']='WEAVE object name from coordinates'
	hdul[2].header.comments['TTYPE4']='Identifier of the target assigned by survey'
	hdul[2].header.comments['TTYPE5']='Input Spectrum normalised'
	hdul[2].header.comments['TTYPE6']='Fit Spectrum normalised'

	hdul[2].header.comments['TFORM1']='data format of field: integer'
	hdul[2].header.comments['TFORM2']='data format of field: integer'
	hdul[2].header.comments['TFORM3']='data format of field: ASCII Character'
	hdul[2].header.comments['TFORM4']='data format of field: ASCII Character'
	hdul[2].header.comments['TFORM5']='data format of field: 4-byte REAL'
	hdul[2].header.comments['TFORM6']='data format of field: 4-byte REAL'

	hdul[2].header.set('TDISP1','I4',after='TFORM1')
	hdul[2].header.set('TDISP2','I4',after='TFORM2')
	hdul[2].header.set('TDISP3','A20',after='TFORM3')
	hdul[2].header.set('TDISP4','A30',after='TFORM4')

	hdul[2].header.comments['TDISP1']='Display format for column'
	hdul[2].header.comments['TDISP2']='Display format for column'
	hdul[2].header.comments['TDISP3']='Display format for column'
	hdul[2].header.comments['TDISP4']='Display format for column'

	hdul[2].header.set('TUCD1','meta.id',after='TDISP1')
	hdul[2].header.set('TUCD2','meta.id',after='TDISP2')
	hdul[2].header.set('TUCD3','meta.id;meta.main',after='TDISP3')
	hdul[2].header.set('TUCD4','meta.id',after='TDISP4')
	hdul[2].header.set('TUCD5','phot.count',after='TFORM5')
	hdul[2].header.set('TUCD6','phot.count',after='TFORM6')
	hdul[2].header.comments['TUCD1']='UCD for column'
	hdul[2].header.comments['TUCD2']='UCD for column'
	hdul[2].header.comments['TUCD3']='UCD for column'
	hdul[2].header.comments['TUCD4']='UCD for column'
	hdul[2].header.comments['TUCD5']='UCD for column'
	hdul[2].header.comments['TUCD6']='UCD for column'

	hdul[2].header.set('TUNIT5','counts','physical unit of field',after='TUCD5')
	hdul[2].header.set('TUNIT6','counts','physical unit of field',after='TUCD6')

	hdul[2].header.set('TDMIN1',1,after='TUCD1')
	hdul[2].header.set('TDMIN2',1,after='TUCD2')
	hdul[2].header.comments['TDMIN1']='Minimum value expected for field'
	hdul[2].header.comments['TDMIN2']='Minimum value expected for field'
	hdul[2].header.set('TDMAX1',960,after='TUCD1')
	hdul[2].header.set('TDMAX2',1100,after='TUCD2')
	hdul[2].header.comments['TDMAX1']='Maximum value expected for field'
	hdul[2].header.comments['TDMAX2']='Maximum value expected for field'
	hdul[2].header.set('TPROP1',0,after='TUCD1')
	hdul[2].header.set('TPROP2',0,after='TUCD2')
	hdul[2].header.set('TPROP3',0,after='TUCD3')
	hdul[2].header.set('TPROP4',0,after='TUCD4')
	hdul[2].header.set('TPROP5',0,after='TUCD5')
	hdul[2].header.set('TPROP6',0,after='TUCD6')
	hdul[2].header.comments['TPROP1']='Public column'
	hdul[2].header.comments['TPROP2']='Public column'
	hdul[2].header.comments['TPROP3']='Public column'
	hdul[2].header.comments['TPROP4']='Public column'
	hdul[2].header.comments['TPROP5']='Public column'
	hdul[2].header.comments['TPROP6']='Public column'

	hdul[2].header.set('TCTYP5','AWAV','Coordinate type',after='TPROP5')
	hdul[2].header.set('TCTYP6','AWAV','Coordinate type',after='TPROP6')
	hdul[2].header.set('TCUNI5','Angstrom','Coordinate unit',after='TCTYP5')
	hdul[2].header.set('TCUNI6','Angstrom','Coordinate unit',after='TCTYP6')
	hdul[2].header.set('TCRPX5',1,'Pixel coordinate of the reference point',after='TCUNI5')
	hdul[2].header.set('TCRPX6',1,'Pixel coordinate of the reference point',after='TCUNI6')
	hdul[2].header.set('TCRVL5',8470.25,'Coordinate value at reference point',after='TCRPX5')
	hdul[2].header.set('TCRVL6',8470.25,'Coordinate value at reference point',after='TCRPX6')
	hdul[2].header.set('TCDLT5',0.25,'Coordinate increment at reference point',after='TCRVL5')
	hdul[2].header.set('TCDLT6',0.25,'Coordinate increment at reference point',after='TCRVL6')

#	hdul[2].header['XTENSION']='BINSPEC'
	return hdul

def createDAT():
	tab = open(write_directory+'results/'+ os.path.splitext(os.path.basename(args.infiles[0]))[0]+'_AMY.dat', "w")
	tab.write("Date of creation: {}\n".format(datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')))
	tab.write("\n")
	tab.write("######## OB info and inputs ###########\n")
	tab.write("OB: {}\n".format(os.path.splitext(os.path.basename(args.infiles[0]))[0].split('_')[1]))
	tab.write("input file: {}\n".format(data_file))
	tab.write("output directory: {}\n".format(write_directory))
	tab.write("output file: {}\n".format(os.path.splitext(os.path.basename(args.infiles[0]))[0]+'_AMY.fits'))
	tab.write("targlist: {}\n".format(target_list))
	tab.write("apsclassification: {}\n".format(apsclassification))
	tab.write("override: {}\n".format(override))

	tab.write("\n")
	tab.write("######## libraries version ###########\n")

	tab.write("numpy version: {}\n".format(np.version.version))
	tab.write("astropy version: {}\n".format(astropy.version.version))
	tab.write("scipy version: {}\n".format(scipy.version.version))
	tab.write("matplotlib version: {}\n".format(matplotlib.__version__))
#	tab.write(": {}\n".format())
	tab.write("\n")
	tab.write("######## parameters.py file: ###############\n")
	tab.close()
	os.system('more parameters.py >>{}'.format(write_directory+'results/'+ os.path.splitext(os.path.basename(args.infiles[0]))[0]+'_AMY.dat'))

def make_output_fits():

	#PHU
	hdr = fits.Header()
	hdr['COMMENT'] = "WEAVE Contributed Software: AMY"
	hdr['DATAMVER'] = 7.60
	hdr.comments['DATAMVER']='WEAVE Data Model Version'
	hdr['CS_CODE'] = 'AMY'
	hdr.comments['CS_CODE']='CS code name'
	hdr['CS_VER'] = 'May20'
	hdr.comments['CS_VER']='CS version '
	hdr['CS_NME1'] = 'Amy, Maria'
	hdr.comments['CS_NME1']='CS author forename'
	hdr['CS_NME2'] = 'Harris, Monguio'
	hdr.comments['CS_NME2']='CS author surname(s)'
	hdr['CS_MAIL'] = 'm.monguio@icc.ub.edu'
	hdr.comments['CS_MAIL']='CS author email'
	hdr['PROV1001'] = os.path.splitext(os.path.basename(args.infiles[0]))[0]+'.fit'
	hdr.comments['PROV1001']='L1 file used'
	hdr['PROV2001'] = ''
	hdr.comments['PROV2001']='L2 file used'
	hdr['DATETIME'] = datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')
	hdr.comments['DATETIME']='Datetime file created'
	print(hdr['PROV1001'],hdr['DATETIME'])
	empty_primary = fits.PrimaryHDU(header=hdr)

	#bin,PARAM
	output_files = glob.glob(write_directory + 'results/*_results')
	all_res = []
	for i in output_files:
		with open(i) as outf:
			all_res.append(outf.read().split())
	t = Table(rows=all_res, names=('Nspec', 'FIBREID', 'CNAME', 'TARGID', 'AMY_ACCEPTANCE', 'AMY_TEFF', 'AMY_TEFF_minus', 'AMY_TEFF_plus', 'AMY_LOGG', 'AMY_LOGG_minus', 'AMY_LOGG_plus', 'AMY_VSINI', 'vsini_minus', 'AMY_VSINI_plus', 'AMY_RV', 'AMY_RV_minus', 'AMY_RV_plus', 'AMY_SLOPE', 'AMY_SLOPE_minus', 'AMY_SLOPE_plus', 'AMY_INTERCEPT', 'AMY_INTERCEPT_minus', 'INTERCEPT_plus'))

	cols=['Nspec', 'FIBREID', 'CNAME', 'TARGID', 'AMY_ACCEPTANCE', 'AMY_TEFF', 'AMY_TEFF_minus','AMY_TEFF_plus', 'AMY_LOGG', 'AMY_LOGG_minus', 'AMY_LOGG_plus', 'AMY_VSINI','vsini_minus', 'AMY_VSINI_plus', 'AMY_RV', 'AMY_RV_minus', 'AMY_RV_plus', 'AMY_SLOPE','AMY_SLOPE_minus', 'AMY_SLOPE_plus', 'AMY_INTERCEPT', 'AMY_INTERCEPT_minus', 'INTERCEPT_plus']
	formats=['I','I','20A','30A','E','E','E','E','E','E','E', 'E','E','E','E','E','E','E','E','E','E','E','E']

	columns=[]
	for i in range(len(cols)):
		columns.append(fits.Column(name=cols[i], array=t[cols[i]],format=formats[i]))

	bintable = fits.BinTableHDU.from_columns(columns)

	#binspec
	output_files = glob.glob(write_directory + 'results/*_spec')
	all_spec = []
	col0=[]
	col1=[]
	col2=[]
	col3=[]
	col4=[]
	col5=[]
	for i in output_files:
		with open(i) as outf:
			line=outf.read().split(';')
			col0.append(line[0])
			col1.append(line[1])
			col2.append(line[2])
			col3.append(line[3])
			col4.append(np.array(np.matrix(line[4])).ravel())
			col5.append(np.array(np.matrix(line[5])).ravel())
	t=[col0,col1,col2,col3,col4,col5]

	cols=['Nspec', 'FIBREID', 'CNAME', 'TARGID','AMY_SPECTRAIN','AMY_SPECTRAFIT']
	formats=['I','I','20A','30A','1879E','1879E']

	columns=[]
	for i in range(len(cols)):
		columns.append(fits.Column(name=cols[i], array=t[i],format=formats[i]))

	spectable = fits.BinTableHDU.from_columns(columns)

	hdul = fits.HDUList([empty_primary,bintable,spectable])
	hdul = modheader(hdul)

	hdul.writeto(write_directory+'results/'+os.path.splitext(os.path.basename(args.infiles[0]))[0]+'_AMY.fits', overwrite=False, checksum=True)

#ascii file with info for reproducibility:
	createDAT()

if __name__ ==  '__main__':
	# checking if results table already exists
	if os.path.exists(write_directory+'results/'+os.path.splitext(os.path.basename(args.infiles[0]))[0]+'_AMY.fits'):
		if override:
			os.remove(write_directory+'results/'+ os.path.splitext(os.path.basename(args.infiles[0]))[0]+'_AMY.fits')
		else:
			print('WARNING: result table already exists, ending process.')
			sys.exit()
	if target_list == 'all':
		print("Processing all BA stars in fits file")
		BA = []
		with fits.open(data_file) as ALLDATA:
			for n,i in enumerate(ALLDATA['FIBTABLE'].data['TARGID']):
				if 'LR-BA' in i:
#			for n,i in enumerate(ALLDATA['FIBTABLE'].data['TARGCLASS']):
#				if 'STAR_BA' in i:
					if ALLDATA['FIBTABLE'].data['STATUS'][n]=='A':
						BA.append(i)
	else:
		print("Processing BA stars specified in target list")
		targlist = np.genfromtxt(target_list, dtype=None, encoding=None)
		BA = []
		with fits.open(data_file) as ALLDATA:
			for targ in targlist:
				try:
					idx=list(ALLDATA['FIBTABLE'].data['FIBREID']).index(targ)
					if ALLDATA['FIBTABLE'].data['STATUS'][idx]=='A':
						BA.append(ALLDATA['FIBTABLE'].data['TARGID'][idx])
				except:
					try:
						idx=list(ALLDATA['FIBTABLE'].data['TARGID']).index(targ)
						if ALLDATA['FIBTABLE'].data['STATUS'][idx]=='A':
							BA.append(ALLDATA['FIBTABLE'].data['TARGID'][idx])
					except:
						print(str(targ)+": Cant find either FIBREID or TARGID in input table.")

ALLDATA = None
info = None
wavelength = None
targetmask = None

def set_globals():
    global ALLDATA
    global info
    global wavelength
    global targetmask

    with fits.open(data_file) as ALLDATA:
	    head0 = ALLDATA[0].header
	    info = ALLDATA['FIBTABLE'].data
	    data1 = ALLDATA['RED_DATA'].data
	    head1 = ALLDATA['RED_DATA'].header
	    wave0 = head1['CRVAL1']
	    increm = head1['CD1_1']
	    wavelength = np.array([wave0+increm*a for a in range(len(data1[1]))])
	    targetmask = (wavelength > min_wav) & (wavelength < max_wav)
	    wavelength = wavelength[targetmask]

if multiprocess==True:
	from multiprocessing import Pool
	if __name__ ==  '__main__':
		pool = Pool(processes=process_no, initializer=set_globals)
		it = pool.imap_unordered(mcmc_one, BA)
		for nn,i in enumerate(range(len(BA))):
			it.next()
		make_output_fits()
		print(datetime.now() - startTime)

else:
	set_globals()
	for nn,i in enumerate(BA):
		mcmc_one(i)
	make_output_fits()
	print(datetime.now() - startTime)

import numpy,import scipy,from scipy,import astropy,from astropy


from det3d.datasets.utils.eval import box3d_overlap
from det3d.core.bbox.box_np_ops import center_to_corner_box3d
import numpy as np

def lvx_track(annos):
    """ 将上一帧的next预测与当前帧对比，将当前帧的prev预测与上一帧对比 """
    new_id = 0
    reserve_id = []
    threshold = -0.125
    threshold_1 = -0.325
    threshold_2 = -0.5

    birth_min = 3
    death_min = 6

    # 上一帧没有匹配的检测框，进入death
    vir_annos = {'track_id':np.array([-2]).reshape(1),
                 'location':np.array([[9999,9999,-2]]).reshape(1,3),
                 'dimensions':np.array([[0,0,0]]).reshape(1,3),
                 'rotation_y':np.array([0]).reshape(1),
                 'location_1':np.array([[9999,9999,-2]]).reshape(1,3),
                 'dimensions_1':np.array([[0,0,0]]).reshape(1,3),
                 'rotation_y_1':np.array([0]).reshape(1),
                 'location_2':np.array([[9999,9999,-2]]).reshape(1,3),
                 'dimensions_2':np.array([[0,0,0]]).reshape(1,3),
                 'rotation_y_2':np.array([0]).reshape(1),
                 'age':np.array([-10000]).reshape(1) # 寿命暂定为6帧
                 }

    # 新生track的id池，连续3帧才算为正式的
    birth_annos = {'track_id':np.array([-2]).reshape(1),
                 'age':np.array([-10000]).reshape(1) # 寿命暂定为6帧
                }

    for i,anno in enumerate(annos):
        if i == 0:
            anno["track_id"] = np.arange(0,anno["name"].shape[0])
            new_id+=anno["name"].shape[0]
        else:
            boxes = anno_to_boxes(anno,'')
            boxes_1 = anno_to_boxes(annos[i-1],'_2')

            boxes_2 = anno_to_boxes(anno,'_1')
            boxes_3 = anno_to_boxes(annos[i-1],'')
            flag = []

            # 先判断是否多个人交汇成一个人
            past_predict = get_track_scores(boxes_2,boxes_3)
            predict_pre = get_track_scores(boxes,boxes_1)

            for age,new_track in zip(birth_annos['age'],birth_annos['track_id']):
                ind = np.where(annos[i-1]['track_id'] == new_track)
                # 假阳性的不进行分裂
                predict_pre[:,ind] = -1
            split = []
            det_match = (predict_pre>threshold_1).sum(1).reshape(-1)
            track_match = (predict_pre>threshold_1).sum(0).reshape(-1)
            for det in range(det_match.shape[0]):
                indices = np.where(predict_pre[det]>threshold_1)[0]
                if det_match[det]>1 and (track_match[indices]<=1).all():
                    # 当前的detect与多个track匹配，说明多个人汇合一起，需要将detect分割
                    # TODO: virtual box应该也算进来

                    # 如果汇合处的detect和其中一个track匹配度非常高，而一些很低，则说明只是个别没有检测到
                    stay = []
                    if predict_pre[det].max()>0.6:
                        for ind in range(len(indices)):
                            if predict_pre[det,indices[ind]] < predict_pre[det].max():
                                stay.append(ind)
                    indices = np.delete(indices,stay)

                    split.append(det)
                    detect = boxes[det].copy()
                    tracks = boxes_1[indices].copy()
                    loc_err = detect-tracks.mean(0)
                    new_detect = tracks.copy()
                    new_detect[:,:3] = tracks[:,:3] + loc_err[:3]
                    new_detect[:,6] = tracks[:,6] + loc_err[6]
                    new_detect_1 = boxes_3[indices].copy()
                    new_detect_2 = new_detect.copy()
                    new_detect_2[:,:3] = new_detect[:,:3]+tracks[:,:3]-new_detect_1[:,:3]
                    boxes = np.concatenate([boxes,new_detect],axis=0)
                    boxes_2 = np.concatenate([boxes_2,new_detect_1],axis=0)
                    anno['location'] = np.concatenate([anno['location'],new_detect[:,:3]],axis=0)
                    anno['dimensions'] = np.concatenate([anno['dimensions'],new_detect[:,3:6]],axis=0)
                    anno['rotation_y'] = np.concatenate([anno['rotation_y'],new_detect[:,6]],axis=0)
                    anno['location_1'] = np.concatenate([anno['location_1'],new_detect_1[:,:3]],axis=0)
                    anno['dimensions_1'] = np.concatenate([anno['dimensions_1'],new_detect_1[:,3:6]],axis=0)
                    anno['rotation_y_1'] = np.concatenate([anno['rotation_y_1'],new_detect_1[:,6]],axis=0)
                    anno['location_2'] = np.concatenate([anno['location_2'],new_detect_2[:,:3]],axis=0)
                    anno['dimensions_2'] = np.concatenate([anno['dimensions_2'],new_detect_2[:,3:6]],axis=0)
                    anno['rotation_y_2'] = np.concatenate([anno['rotation_y_2'],new_detect_2[:,6]],axis=0)
                    anno['name'] = np.concatenate([anno['name'],annos[i-1]['name'][indices]],axis=0)
                    anno['alpha'] = np.concatenate([anno['alpha'],annos[i-1]['alpha'][indices]],axis=0)
                    anno['bbox'] = np.concatenate([anno['bbox'],annos[i-1]['bbox'][indices]],axis=0)
                    anno['score'] = np.concatenate([anno['score'],anno['score'][det:det+1].repeat(indices.shape[0],axis=0)],axis=0)

            boxes = np.delete(boxes,split,axis=0)
            boxes_2 = np.delete(boxes_2,split,axis=0)
            track_ids = -1*np.ones(boxes.shape[0])
            scores = -1*np.ones(boxes.shape[0])

            # 将detect分割后同时修改anno
            anno = del_dict(anno,split)

            assert annos[i]['name'].shape[0] == boxes.shape[0]

            # 基准帧为过去一帧的当前检测
            past_predict = get_track_scores(boxes_2,boxes_3)
            # get_track(past_predict,threshold,track_ids,scores, annos[i-1],flag)

            predict_pre = get_track_scores(boxes,boxes_1)

            # T,T-1的可信度更高，T+1的可信度更低
            predict_pre = 0.6*past_predict+0.4*predict_pre
            for age,new_track in zip(birth_annos['age'],birth_annos['track_id']):
                ind = np.where(annos[i-1]['track_id'] == new_track)
                # new track不确定是否为假阳性，所以权重调低
                predict_pre[:,ind] = (predict_pre[:,ind]-threshold)*(1+age)/(birth_min+1) + threshold
            get_track(predict_pre,threshold,track_ids,scores,annos[i-1],flag)

            # 前一帧没有被匹配的加入到vir中，并且是已经确定为非假阳性的，不在birth里面
            for j in range(len(annos[i-1]['track_id'])):
                if j not in flag and annos[i-1]['track_id'][j] not in birth_annos['track_id']:
                    # 否则，则将其添加到vir_annos里
                    vir_annos['track_id']=np.concatenate([vir_annos['track_id'],[annos[i-1]['track_id'][j]]],axis=-1)
                    vir_annos['location']=np.concatenate([vir_annos['location'],[annos[i-1]['location'][j]]],axis=0)
                    vir_annos['dimensions']=np.concatenate([vir_annos['dimensions'],[annos[i-1]['dimensions'][j]]],axis=0)
                    vir_annos['rotation_y']=np.concatenate([vir_annos['rotation_y'],[annos[i-1]['rotation_y'][j]]],axis=-1)
                    vir_annos['location_1']=np.concatenate([vir_annos['location_1'],[annos[i-1]['location_1'][j]]],axis=0)
                    vir_annos['dimensions_1']=np.concatenate([vir_annos['dimensions_1'],[annos[i-1]['dimensions_1'][j]]],axis=0)
                    vir_annos['rotation_y_1']=np.concatenate([vir_annos['rotation_y_1'],[annos[i-1]['rotation_y_1'][j]]],axis=-1)
                    vir_annos['location_2']=np.concatenate([vir_annos['location_2'],[annos[i-1]['location_2'][j]]],axis=0)
                    vir_annos['dimensions_2']=np.concatenate([vir_annos['dimensions_2'],[annos[i-1]['dimensions_2'][j]]],axis=0)
                    vir_annos['rotation_y_2']=np.concatenate([vir_annos['rotation_y_2'],[annos[i-1]['rotation_y_2'][j]]],axis=-1)
                    vir_annos['age']=np.concatenate([vir_annos['age'],[0]],axis=-1)

            vir_boxes = anno_to_boxes(vir_annos,'')
            flag_vir = []
            # 与virtual框进行阈值更低的匹配
            # if anno['metadata']['token'] == '642':
            #     print(boxes_2,vir_boxes)
            #     import pdb; pdb.set_trace()
            virtual = get_track_scores(boxes_2,vir_boxes)
            get_track(virtual,threshold_2,track_ids,scores,vir_annos,flag_vir)

            # 删除virtual中被成功匹配的框,被具现化了
            vir_annos = del_dict(vir_annos,flag_vir)

            # 向下一帧移动
            vir_annos['location_1'][:,:2]= vir_annos['location'][:,:2].copy()
            vir_annos['location'][:,:2] = vir_annos['location_2'][:,:2].copy()
            flag_vir = []
            for j in range((len(vir_annos['track_id']))):
                vir_annos['location_2'][j,:2] = 2*vir_annos['location'][j,:2]-vir_annos['location_1'][j,:2]
                vir_annos['age'][j] += 1
                if vir_annos['location'][j,0]<-40 or vir_annos['location'][j,0]>40 or vir_annos['location'][j,1]<-40 or vir_annos['location'][j,1]>40:
                    # 如果出了边界，去掉这个人
                    flag_vir.append(j)
                if vir_annos['age'][j] > death_min and j not in flag_vir:
                    # age>6 去掉
                    flag_vir.append(j)

            vir_annos = del_dict(vir_annos,flag_vir)

            #处理新生id
            flag_birth = []
            for j,new_track in enumerate(birth_annos['track_id']):
                ind = np.where(annos[i-1]['track_id'][flag] == new_track)[0]
                if ind.size>0 and birth_annos['age'][j]<birth_min-1:
                    birth_annos['age'][j] += 1
                elif new_track>=0:
                    # 若达到成年寿命，或者中途夭折，去掉
                    flag_birth.append(j)
                    if birth_annos['age'][j] < birth_min-1:
                        reserve_id.append(new_track)
            birth_annos['age'] = np.delete(birth_annos['age'],flag_birth)
            birth_annos['track_id'] = np.delete(birth_annos['track_id'],flag_birth)

            for index in range(track_ids.shape[0]):
                if track_ids[index] == -1:
                    print(anno['metadata']['token'],index,new_id,anno['location'][index])
                    if len(reserve_id) == 0:
                        # 假阳性的id再次使用
                        track_ids[index] = new_id
                        new_id += 1
                    else:
                        track_ids[index] = reserve_id[-1]
                        del reserve_id[-1]
                    birth_annos['track_id'] = np.concatenate([birth_annos['track_id'],[track_ids[index]]],axis=0)
                    birth_annos['age'] = np.concatenate([birth_annos['age'],[0]],axis=0)
            anno["track_id"] = track_ids

    return annos

def get_track_scores(boxes,boxes_1):
    ''' DIoU'''
    predict_pre = box3d_overlap(boxes,boxes_1,z_axis=2,z_center=0.5)

    corners = center_to_corner_box3d(boxes[:,:3],boxes[:,3:6],boxes[:,6])
    corners_1 = center_to_corner_box3d(boxes_1[:,:3],boxes_1[:,3:6],boxes_1[:,6])
    for i in range(boxes.shape[0]):
        for j in range(boxes_1.shape[0]):
            c = ((boxes[i,:3]-boxes_1[j,:3])**2).sum()
            d = get_farest(corners[i],corners_1[j])
            predict_pre[i,j] -= c/d
    return predict_pre

def get_farest(corners,corners_1):
    max_d = 0
    for i in range(8):
        for j in range(8):
            dist = ((corners[i]-corners_1[j])**2).sum()
            if dist > max_d:
                max_d = dist
    return max_d

def anno_to_boxes(anno,flag):
    boxes = np.concatenate([anno[f'location{flag}'],anno[f'dimensions{flag}'],anno[f'rotation_y{flag}'][...,np.newaxis]],axis=1)
    return boxes

def get_track(iou, threshold, track, scores, anno, flag):
    num1 = 0
    num2 = 0
    while(num1<iou.shape[0] and num2<iou.shape[1]):
        loc_max = iou.argmax()
        max_iou = [loc_max//iou.shape[1],loc_max%iou.shape[1]]
        score = iou[max_iou[0],max_iou[1]]
        if score <threshold:
            break
        # track_ids 需要没有被占用过,并且前一帧的flag没被用过
        if track[max_iou[0]] != -1:
            iou[max_iou[0],:] = -1
            num1+=1
        elif max_iou[1] in flag:
            iou[:,max_iou[1]] = -1
            num2+=1
        else:
            track[max_iou[0]] = anno['track_id'][max_iou[1]]
            scores[max_iou[0]] = score
            flag.append(max_iou[1])
            iou[max_iou[0],:] = -1
            iou[:,max_iou[1]] = -1
            num1+=1
            num2+=1

def del_dict(anno,flag):
    anno['location'] = np.delete(anno['location'],flag,axis=0)
    anno['dimensions'] = np.delete(anno['dimensions'],flag,axis=0)
    anno['rotation_y'] = np.delete(anno['rotation_y'],flag,axis=0)
    anno['location_1'] = np.delete(anno['location_1'],flag,axis=0)
    anno['dimensions_1'] = np.delete(anno['dimensions_1'],flag,axis=0)
    anno['rotation_y_1'] = np.delete(anno['rotation_y_1'],flag,axis=0)
    anno['location_2'] = np.delete(anno['location_2'],flag,axis=0)
    anno['dimensions_2'] = np.delete(anno['dimensions_2'],flag,axis=0)
    anno['rotation_y_2'] = np.delete(anno['rotation_y_2'],flag,axis=0)

    if 'name' in anno.keys():
        anno['name'] = np.delete(anno['name'],flag,axis=0)
        anno['alpha'] = np.delete(anno['alpha'],flag,axis=0)
        anno['score'] = np.delete(anno['score'],flag,axis=0)
        anno['bbox'] = np.delete(anno['bbox'],flag,axis=0)

    if 'age' in anno.keys():
        anno['track_id'] = np.delete(anno['track_id'],flag,axis=0)
        anno['age'] = np.delete(anno['age'],flag,axis=0)

    return anno


import numpy as np
import pandas as pd
import pymc3 as pm
from abc import ABC, abstractmethod
from dataclasses import dataclass
from .constants import *
from .utils import *

__all__ = ['Damuta', 'DataSet', 'SignatureSet']

_opt_methods = {"ADVI": pm.ADVI, "FullRankADVI": pm.FullRankADVI}
_init_strats = ['kmeans', 'uniform']

@dataclass
class DataSet:
    """Container for tabular data, allowing simple access to a mutation data set and corresponding annotation for each sample.

    :class:`DataSet` is instatiated from a pandas dataframe of mutation counts, and (optionally) a pandas dataframe of the
    same size of sample annotations. The dataframe index is taken as sample ids. All samples that appear in counts should
    also appear in annotation, and vice versa. Mutation types are expect to be in COSMIC format (ex. A[C>A]A).

    Parameters
    ----------
    counts: pd.DataFrame
        Nx96 dataframe of mutation counts, one sample per row. Index is assumed to be sample ids.
    annotation: pd.DataFrame
        NxF dataframe of meta-data features to annotate samples with. Index is assumed to be sample ids.

    Examples
    --------
    >>> import pandas as pd
    >>> counts = pd.read_csv('tests/test_data/pcawg_counts.csv', index_col = 0, header = 0)
    >>> annotation = pd.read_csv('tests/test_data/pcawg_cancer_types.csv', index_col = 0, header = 0)
    >>> pcawg = DataSet(counts, annotation)
    >>> pcawg.nsamples
    2778
    """

    counts: pd.DataFrame
    annotation: pd.DataFrame = None

    def __post_init__(self):
        if self.counts is not None:
            assert self.counts.ndim == 2, f'Expected counts.ndim==2. Got {self.counts.ndim}'
            assert self.counts.shape[1] == 96, f'Expected 96 mutation types, got {self.counts.shape[1]}'
            assert all(self.counts.columns.isin(mut96)), 'Unexpected mutation type. Check the counts.columns are in COSMIC mutation type format (ex. A[C>A]A). See COSMIC database for more.'
            # reorder columns if necessary
            self.counts = self.counts[mut96]

        if self.annotation is not None:
            # check the counts and annotation match
            assert self.annotation.shape[0] == self.counts.shape[0], f"Shape mismatch. Expected self.annotation.shape[0] == self.counts.shape[0], got {self.annotation.shape[0]}, {self.counts.shape[0]}"
            assert self.annotation.index.isin(self.counts.index).all() and self.counts.index.isin(self.annotation.index).all(), "Counts and annotation indices must match"

    @property
    def n_samples(self) -> int:
        """Number of samples in dataset"""
        return self.counts.shape[0]

    @property
    def ids(self) -> list:
        """List sample ids in dataset"""
        return self.counts.index.to_list()

    def annotate_tissue_types(self, type_col) -> np.array:
        """Set a specified column of annotation as the sample tissue type

        Tissue type information is used by hirearchical models to create tissue-type prior.
        See class:`HierarchicalTendemLda` for more details.
        """
        if self.annotation is None:
            raise ValueError('Dataset annotation must be provided.')
        assert type_col in self.annotation.columns, f"{type_col} not found in annotation columns. Check spelling?"
        self.tissue_types = pd.Categorical(self.annotation[type_col])
        self.type_codes = self.tissue_types.codes

@dataclass
class SignatureSet:
    """Container for tabular data, allowing simple access to a set of mutational signature definitions.

    Parameters
    ----------
    signatures: pd.DataFrame
        Nx96 dataframe of signautre definitions, one signature per row. Rows must sum to 1.

    Examples
    ----------
    """

    signatures: pd.DataFrame

    def __post_init__(self):
        # check for shape, valid signautre definitions
        assert self.signatures.shape[1] == 96, f"Expected 96 mutation types, got {self.signatures.shape[1]}"
        assert np.allclose(self.signatures.sum(1),1), "All signature definitions must sum to 1"

    @property
    def n_sigs(self) -> int:
        """Number of signatures in dataset"""
        return self.signatures.shape[0]

    @property
    def damage_signatures(self) -> pd.DataFrame:
        """Damage signatures

        Damage signatures represent the distribution of mutations over 32 trinucleotide contexts.
        They are computed by marginalizing over substitution classes.
        """
        phi = get_phi(self.signatures.to_numpy())
        return pd.DataFrame(phi, index = self.signatures.index, columns=mut32)

    @property
    def misrepair_signatures(self) -> pd.DataFrame:
        """Misrepair signatures

        Misrepair signatures represent the distribution of mutations over 6 substitution types.
        They are computed by marginalizing over trinucleotide context classes.
        """
        eta = get_eta(self.signatures.to_numpy())
        return pd.DataFrame(eta, index = self.signatures.index, columns=mut6)

    def summarize_separation(self) -> pd.DataFrame:
        """Summary statistics of pair-wise cosine distances for signautres,
        damage signatures, and misrepair signatures.

        """

        seps = {'Signature separation': cosine_similarity(self.signatures)[np.triu_indices(self.n_sigs, k=1)],
                'Damage signature separation': cosine_similarity(self.damage_signatures)[np.triu_indices(self.n_sigs, k=1)],
                'Misrepair signature separation': cosine_similarity(self.misrepair_signatures)[np.triu_indices(self.n_sigs, k=1)]
               }

        return pd.DataFrame.from_dict(seps).describe()

class Model(ABC):
    """
    Bayesian inference of mutational signautres and their activities.

    The Damuta class acts as a central interface for several types of latent models. Each subclass defines at least `build_model`,
    `fit`, `predict_activities`, `model_to_gv` and metrics such as `LAP`, `ALP`, and `BOR` in addition to subclass-specific methods.

    Parameters
    ----------
    dataset : DataSet
        Data for fitting.
    opt_method: str
        one of "ADVI" for mean field inference, or "FullRankADVI" for full rank inference.
    seed : int
        Random seed

    Attributes
    ----------
    model: pymc3.model.Model object
        pymc3 model instance
    approx: pymc3.variational.approximations object
        pymc3 approximation object. Created via self.fit()
     """

    def __init__(self, dataset: DataSet, opt_method: str, init_strategy: str, seed: int):

        if not isinstance(dataset, DataSet):
            raise TypeError('Learner instance must be initialized with a DataSet object')

        if not opt_method in _opt_methods.keys():
            raise TypeError(f'Optimization method should be one of {list(_opt_methods.keys())}')
        assert init_strategy in _init_strats, f'self.init_strategy should be one of {_init_strats}'

        self.dataset = dataset
        self.opt_method = opt_method
        self.init_strategy = init_strategy
        self.seed = seed
        self.model = None
        self.approx = None

        # hidden attributes
        self._model_kwargs = None
        self._opt = _opt_methods[self.opt_method]
        self._trace = None
        self._hat = None
        self._rng = np.random.default_rng(self.seed)

        # set seed
        np.random.seed(self.seed)
        pm.set_tt_rng(self.seed)

    ################################################################################
    # Model building and fitting
    ################################################################################

    @abstractmethod
    def _build_model(self, *args, **kwargs):
        """Build the pymc3 model
        """
        pass

    @abstractmethod
    def _init_kmeans(self):
        """Defined by subclass
        """
        pass

    @abstractmethod
    def _initialize_signatures(self):
        """Defined by subclass.
        """
        pass

    def fit(self, n, **pymc3_kwargs):
        """Fit model to the dataset specified by self.dataset

        Parameters
        ----------
        n: int
            Number of iterations
        **pymc3_kwargs:
            More parameters to pass to pymc3.fit() (ex. callbacks)

        Returns
        -------
        self: :class:`Lda`
        """

        self._initialize_signatures()
        self._build_model(**self._model_kwargs)

        with self.model:
            self._trace = self._opt(random_seed = self.seed)
            self._trace.fit(n=n, **pymc3_kwargs)

        self.approx = self._trace.approx

        return self

    #@abstractmethod
    #def predict_activites(self, new_data, *args, **kwargs):
    #    """Defined by subclass
    #    """
    #    pass

    #def model_to_gv(self, *args, **kwargs):
    #    """Defined by subclass
    #    """
    #    pass

    ################################################################################
    # Metrics
    ################################################################################

    @abstractmethod
    def BOR(self, *args, **kwargs):
        """Defined by subclass
        """
        pass

    def ALP(self, n_samples = 20):
        """Average log probability per mutation
        """
        if self.approx is None:
            warnings.warn("self.approx is None... Fit the model first!", ValueError)

        B = self.approx.sample(n_samples).B.mean(0)
        return alp_B(self.dataset.counts.to_numpy(), B)

    def LAP(self, n_samples = 20):
        """Log average data likelihood (bayesian version of reconstruction error)
        """
        if self.approx is None:
            warnings.warn("self.trace is None... Fit the model first!", ValueError)

        B = self.approx.sample(n_samples).B.mean(0)
        return alp_B(self.dataset.counts.to_numpy(), B)

    def reconstruction_err(self, *args, **kwargs):
        """Defined by subclass
        """
        pass


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__author__ = "Christian Heider Nielsen"
__doc__ = r"""

           Created on 22/03/2020
           """

import itertools
import logging
import os
from collections import defaultdict
from datetime import datetime
from typing import Tuple

import numpy
import six

__all__ = [
    "bbox_iou",
    "eval_detection_voc",
    "calc_detection_voc_ap",
    "calc_detection_voc_prec_rec",
    "voc_evaluation",
]

def bbox_iou(bbox_a: numpy.ndarray, bbox_b: numpy.ndarray) -> numpy.ndarray:
    """Calculate the Intersection of Unions (IoUs) between bounding boxes.
IoU is calculated as a ratio of area of the intersection
and area of the union.
This function accepts both :obj:`numpy.ndarray` and :obj:`cupy.ndarray` as
inputs. Please note that both :obj:`bbox_a` and :obj:`bbox_b` need to be
same type.
The output is same type as the type of the inputs.
Args:
bbox_a (array): An array whose shape is :math:`(N, 4)`.
  :math:`N` is the number of bounding boxes.
  The dtype should be :obj:`numpy.float32`.
bbox_b (array): An array similar to :obj:`bbox_a`,
  whose shape is :math:`(K, 4)`.
  The dtype should be :obj:`numpy.float32`.
Returns:
array:
An array whose shape is :math:`(N, K)`. \
An element at index :math:`(n, k)` contains IoUs between \
:math:`n` th bounding box in :obj:`bbox_a` and :math:`k` th bounding \
box in :obj:`bbox_b`.
"""
    if bbox_a.shape[1] != 4 or bbox_b.shape[1] != 4:
        raise IndexError

    # top left
    tl = numpy.maximum(bbox_a[:, None, :2], bbox_b[:, :2])
    # bottom right
    br = numpy.minimum(bbox_a[:, None, 2:], bbox_b[:, 2:])

    area_i = numpy.prod(br - tl, axis=2) * (tl < br).all(axis=2)
    area_a = numpy.prod(bbox_a[:, 2:] - bbox_a[:, :2], axis=1)
    area_b = numpy.prod(bbox_b[:, 2:] - bbox_b[:, :2], axis=1)
    return area_i / (area_a[:, None] + area_b - area_i)

def eval_detection_voc(
    pred_bboxes,
    pred_labels,
    pred_scores,
    gt_bboxes,
    gt_labels,
    gt_difficults=None,
    iou_thresh=0.5,
    use_07_metric=False,
) -> Tuple:
    """Calculate average precisions based on evaluation code of PASCAL VOC.

This function evaluates predicted bounding boxes obtained from a dataset
which has :math:`N` images by using average precision for each class.
The code is based on the evaluation code used in PASCAL VOC Challenge.

Args:
pred_bboxes (iterable of numpy.ndarray): An iterable of :math:`N`
  sets of bounding boxes.
  Its index corresponds to an index for the base dataset.
  Each element of :obj:`pred_bboxes` is a set of coordinates
  of bounding boxes. This is an array whose shape is :math:`(R, 4)`,
  where :math:`R` corresponds
  to the number of bounding boxes, which may vary among boxes.
  The second axis corresponds to
  :math:`y_{min}, x_{min}, y_{max}, x_{max}` of a bounding box.
pred_labels (iterable of numpy.ndarray): An iterable of labels.
  Similar to :obj:`pred_bboxes`, its index corresponds to an
  index for the base dataset. Its length is :math:`N`.
pred_scores (iterable of numpy.ndarray): An iterable of confidence
  scores for predicted bounding boxes. Similar to :obj:`pred_bboxes`,
  its index corresponds to an index for the base dataset.
  Its length is :math:`N`.
gt_bboxes (iterable of numpy.ndarray): An iterable of ground truth
  bounding boxes
  whose length is :math:`N`. An element of :obj:`gt_bboxes` is a
  bounding box whose shape is :math:`(R, 4)`. Note that the number of
  bounding boxes in each image does not need to be same as the number
  of corresponding predicted boxes.
gt_labels (iterable of numpy.ndarray): An iterable of ground truth
  labels which are organized similarly to :obj:`gt_bboxes`.
gt_difficults (iterable of numpy.ndarray): An iterable of boolean
  arrays which is organized similarly to :obj:`gt_bboxes`.
  This tells whether the
  corresponding ground truth bounding box is difficult or not.
  By default, this is :obj:`None`. In that case, this function
  considers all bounding boxes to be not difficult.
iou_thresh (float): A prediction is correct if its Intersection over
  Union with the ground truth is above this value.
use_07_metric (bool): Whether to use PASCAL VOC 2007 evaluation metric
  for calculating average precision. The default value is
  :obj:`False`.

Returns:
dict:

The keys, value-types and the description of the values are listed
below.

* **ap** (*numpy.ndarray*): An array of average precisions. \
  The :math:`l`-th value corresponds to the average precision \
  for class :math:`l`. If class :math:`l` does not exist in \
  either :obj:`pred_labels` or :obj:`gt_labels`, the corresponding \
  value is set to :obj:`numpy.nan`.
* **map** (*float*): The mean of Average Precisions over classes.

"""

    prec, rec = calc_detection_voc_prec_rec(
        pred_bboxes,
        pred_labels,
        pred_scores,
        gt_bboxes,
        gt_labels,
        gt_difficults,
        iou_thresh=iou_thresh,
    )

    ap = calc_detection_voc_ap(prec, rec, use_07_metric=use_07_metric)

    return ap, numpy.nanmean(ap)  # Mean Average Precision

def calc_detection_voc_prec_rec(
    pred_bboxes,
    pred_labels,
    pred_scores,
    gt_bboxes,
    gt_labels,
    gt_difficults=None,
    iou_thresh: float = 0.5,
) -> Tuple:
    """Calculate precision and recall based on evaluation code of PASCAL VOC.

This function calculates precision and recall of
predicted bounding boxes obtained from a dataset which has :math:`N`
images.
The code is based on the evaluation code used in PASCAL VOC Challenge.

Args:
pred_bboxes (iterable of numpy.ndarray): An iterable of :math:`N`
  sets of bounding boxes.
  Its index corresponds to an index for the base dataset.
  Each element of :obj:`pred_bboxes` is a set of coordinates
  of bounding boxes. This is an array whose shape is :math:`(R, 4)`,
  where :math:`R` corresponds
  to the number of bounding boxes, which may vary among boxes.
  The second axis corresponds to
  :math:`y_{min}, x_{min}, y_{max}, x_{max}` of a bounding box.
pred_labels (iterable of numpy.ndarray): An iterable of labels.
  Similar to :obj:`pred_bboxes`, its index corresponds to an
  index for the base dataset. Its length is :math:`N`.
pred_scores (iterable of numpy.ndarray): An iterable of confidence
  scores for predicted bounding boxes. Similar to :obj:`pred_bboxes`,
  its index corresponds to an index for the base dataset.
  Its length is :math:`N`.
gt_bboxes (iterable of numpy.ndarray): An iterable of ground truth
  bounding boxes
  whose length is :math:`N`. An element of :obj:`gt_bboxes` is a
  bounding box whose shape is :math:`(R, 4)`. Note that the number of
  bounding boxes in each image does not need to be same as the number
  of corresponding predicted boxes.
gt_labels (iterable of numpy.ndarray): An iterable of ground truth
  labels which are organized similarly to :obj:`gt_bboxes`.
gt_difficults (iterable of numpy.ndarray): An iterable of boolean
  arrays which is organized similarly to :obj:`gt_bboxes`.
  This tells whether the
  corresponding ground truth bounding box is difficult or not.
  By default, this is :obj:`None`. In that case, this function
  considers all bounding boxes to be not difficult.
iou_thresh (float): A prediction is correct if its Intersection over
  Union with the ground truth is above this value..

Returns:
tuple of two lists:
This function returns two lists: :obj:`prec` and :obj:`rec`.

* :obj:`prec`: A list of arrays. :obj:`prec[l]` is precision \
  for class :math:`l`. If class :math:`l` does not exist in \
  either :obj:`pred_labels` or :obj:`gt_labels`, :obj:`prec[l]` is \
  set to :obj:`None`.
* :obj:`rec`: A list of arrays. :obj:`rec[l]` is recall \
  for class :math:`l`. If class :math:`l` that is not marked as \
  difficult does not exist in \
  :obj:`gt_labels`, :obj:`rec[l]` is \
  set to :obj:`None`.

"""

    pred_bboxes = iter(pred_bboxes)
    pred_labels = iter(pred_labels)
    pred_scores = iter(pred_scores)
    gt_bboxes = iter(gt_bboxes)
    gt_labels = iter(gt_labels)
    if gt_difficults is None:
        gt_difficults = itertools.repeat(None)
    else:
        gt_difficults = iter(gt_difficults)

    n_pos = defaultdict(int)
    score = defaultdict(list)
    match = defaultdict(list)

    for (
        pred_bbox,
        pred_label,
        pred_score,
        gt_bbox,
        gt_label,
        gt_difficult,
    ) in six.moves.zip(
        pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels, gt_difficults
    ):

        if gt_difficult is None:
            gt_difficult = numpy.zeros(gt_bbox.shape[0], dtype=bool)

        for l in numpy.unique(numpy.concatenate((pred_label, gt_label)).astype(int)):
            pred_mask_l = pred_label == l
            pred_bbox_l = pred_bbox[pred_mask_l]
            pred_score_l = pred_score[pred_mask_l]
            # sort by score
            order = pred_score_l.argsort()[::-1]
            pred_bbox_l = pred_bbox_l[order]
            pred_score_l = pred_score_l[order]

            gt_mask_l = gt_label == l
            gt_bbox_l = gt_bbox[gt_mask_l]
            gt_difficult_l = gt_difficult[gt_mask_l]

            n_pos[l] += numpy.logical_not(gt_difficult_l).sum()
            score[l].extend(pred_score_l)

            if len(pred_bbox_l) == 0:
                continue
            if len(gt_bbox_l) == 0:
                match[l].extend((0,) * pred_bbox_l.shape[0])
                continue

            # VOC evaluation follows integer typed bounding boxes.
            pred_bbox_l = pred_bbox_l.copy()
            pred_bbox_l[:, 2:] += 1
            gt_bbox_l = gt_bbox_l.copy()
            gt_bbox_l[:, 2:] += 1

            iou = bbox_iou(pred_bbox_l, gt_bbox_l)
            gt_index = iou.argmax(axis=1)
            # set -1 if there is no matching ground truth
            gt_index[iou.max(axis=1) < iou_thresh] = -1
            del iou

            selec = numpy.zeros(gt_bbox_l.shape[0], dtype=bool)
            for gt_idx in gt_index:
                if gt_idx >= 0:
                    if gt_difficult_l[gt_idx]:
                        match[l].append(-1)
                    else:
                        if not selec[gt_idx]:
                            match[l].append(1)
                        else:
                            match[l].append(0)
                    selec[gt_idx] = True
                else:
                    match[l].append(0)

    for iter_ in (
        pred_bboxes,
        pred_labels,
        pred_scores,
        gt_bboxes,
        gt_labels,
        gt_difficults,
    ):
        if next(iter_, None) is not None:
            raise ValueError("Length of input iterables need to be same.")

    n_fg_class = max(n_pos.keys()) + 1
    prec = [None] * n_fg_class
    rec = [None] * n_fg_class

    for l in n_pos.keys():
        score_l = numpy.array(score[l])
        match_l = numpy.array(match[l], dtype=numpy.int8)

        order = score_l.argsort()[::-1]
        match_l = match_l[order]

        tp = numpy.cumsum(match_l == 1)
        fp = numpy.cumsum(match_l == 0)

        # If an element of fp + tp is 0,
        # the corresponding element of prec[l] is nan.
        prec[l] = tp / (fp + tp)
        # If n_pos[l] is 0, rec[l] is None.
        if n_pos[l] > 0:
            rec[l] = tp / n_pos[l]

    return prec, rec

def calc_detection_voc_ap(prec, rec, use_07_metric=False):
    """Calculate average precisions based on evaluation code of PASCAL VOC.

This function calculates average precisions
from given precisions and recalls.
The code is based on the evaluation code used in PASCAL VOC Challenge.

Args:
prec (list of numpy.array): A list of arrays.
  :obj:`prec[l]` indicates precision for class :math:`l`.
  If :obj:`prec[l]` is :obj:`None`, this function returns
  :obj:`numpy.nan` for class :math:`l`.
rec (list of numpy.array): A list of arrays.
  :obj:`rec[l]` indicates recall for class :math:`l`.
  If :obj:`rec[l]` is :obj:`None`, this function returns
  :obj:`numpy.nan` for class :math:`l`.
use_07_metric (bool): Whether to use PASCAL VOC 2007 evaluation metric
  for calculating average precision. The default value is
  :obj:`False`.

Returns:
~numpy.ndarray:
This function returns an array of average precisions.
The :math:`l`-th value corresponds to the average precision
for class :math:`l`. If :obj:`prec[l]` or :obj:`rec[l]` is
:obj:`None`, the corresponding value is set to :obj:`numpy.nan`.

"""

    n_fg_class = len(prec)
    ap = numpy.empty(n_fg_class)
    for l in six.moves.range(n_fg_class):
        if prec[l] is None or rec[l] is None:
            ap[l] = numpy.nan
            continue

        if use_07_metric:
            # 11 point metric
            ap[l] = 0
            for t in numpy.arange(0.0, 1.1, 0.1):
                if numpy.sum(rec[l] >= t) == 0:
                    p = 0
                else:
                    p = numpy.max(numpy.nan_to_num(prec[l])[rec[l] >= t])
                ap[l] += p / 11
        else:
            # correct AP calculation
            # first append sentinel values at the end
            mpre = numpy.concatenate(([0], numpy.nan_to_num(prec[l]), [0]))
            mrec = numpy.concatenate(([0], rec[l], [1]))

            mpre = numpy.maximum.accumulate(mpre[::-1])[::-1]

            # to calculate area under PR curve, look for points
            # where X axis (recall) changes value
            i = numpy.where(mrec[1:] != mrec[:-1])[0]

            # and sum (\Delta recall) * prec
            ap[l] = numpy.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])

    return ap

def voc_evaluation(dataset, predictions, output_dir, iteration=None):
    class_names = dataset.class_names

    pred_boxes_list = []
    pred_labels_list = []
    pred_scores_list = []
    gt_boxes_list = []
    gt_labels_list = []
    gt_difficults = []

    for i in range(len(dataset)):
        image_id, annotation = dataset.get_annotation(i)
        gt_boxes, gt_labels, is_difficult = annotation
        gt_boxes_list.append(gt_boxes)
        gt_labels_list.append(gt_labels)
        gt_difficults.append(is_difficult.astype(numpy.bool))

        img_info = dataset.get_img_info(i)
        prediction = predictions[i]
        prediction = prediction.resize((img_info["width"], img_info["height"])).numpy()
        boxes, labels, scores = (
            prediction["boxes"],
            prediction["labels"],
            prediction["scores"],
        )

        pred_boxes_list.append(boxes)
        pred_labels_list.append(labels)
        pred_scores_list.append(scores)
    ap, map = eval_detection_voc(
        pred_bboxes=pred_boxes_list,
        pred_labels=pred_labels_list,
        pred_scores=pred_scores_list,
        gt_bboxes=gt_boxes_list,
        gt_labels=gt_labels_list,
        gt_difficults=gt_difficults,
        iou_thresh=0.5,
        use_07_metric=True,
    )
    logger = logging.getLogger("SSD.inference")
    result_str = f"mAP: {map:.4f}\n"
    metrics = {"mAP": map}
    for i, ap in enumerate(ap):
        if i == 0:  # skip background
            continue
        metrics[class_names[i]] = ap
        result_str += f"{class_names[i]:<16}: {ap:.4f}\n"
    logger.info(result_str)

    if iteration is not None:
        result_path = os.path.join(output_dir, f"result_{iteration:07d}.txt")
    else:
        result_path = os.path.join(
            output_dir, f'result_{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.txt'
        )
    with open(result_path, "w") as f:
        f.write(result_str)

    return dict(metrics=metrics)