import unittest

import numpy as np
import pytest
import torch
from numpy.testing import assert_almost_equal

from torch_audiomentations.augmentations.gain import Gain
from torch_audiomentations.utils.dsp import (
    convert_decibels_to_amplitude_ratio,
    convert_amplitude_ratio_to_decibels,
)

class TestGain(unittest.TestCase):
    def test_gain(self):
        samples = np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-6.000001, max_gain_in_db=-6, p=1.0)
        processed_samples = augment(
            samples=torch.from_numpy(samples), sample_rate=sample_rate
        ).numpy()
        expected_factor = convert_decibels_to_amplitude_ratio(-6)
        assert_almost_equal(
            processed_samples,
            expected_factor
            * np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32),
            decimal=6,
        )
        self.assertEqual(processed_samples.dtype, np.float32)

    def test_gain_per_channel(self):
        samples = np.array(
            [[1.0, 0.5, 0.25, 0.125, 0.01], [0.95, 0.5, 0.25, 0.12, 0.011]],
            dtype=np.float32,
        )
        samples_batch = np.stack([samples] * 1000, axis=0)
        sample_rate = 16000

        augment = Gain(
            min_gain_in_db=-30.0, max_gain_in_db=0.0, mode="per_channel", p=0.5
        )
        processed_samples = augment(
            samples=torch.from_numpy(samples_batch), sample_rate=sample_rate
        ).numpy()

        num_unprocessed_channels = 0
        num_processed_channels = 0
        perturbation_type_counter = {
            "zero_channels_changed": 0,
            "only_the_first_channel_changed": 0,
            "only_the_second_channel_changed": 0,
            "both_channels_changed_differently": 0,
            "both_channels_changed_equally": 0,
        }
        for i in range(processed_samples.shape[0]):
            num_perturbed_channels = 0
            est_gain_factors = []
            for chn_index in range(processed_samples.shape[1]):
                if np.allclose(
                    processed_samples[i, chn_index], samples_batch[i, chn_index]
                ):
                    num_unprocessed_channels += 1
                else:
                    num_processed_channels += 1
                    num_perturbed_channels += 1

                    estimated_gain_factors = (
                        processed_samples[i, chn_index] / samples_batch[i, chn_index]
                    )
                    self.assertAlmostEqual(
                        np.amin(estimated_gain_factors),
                        np.amax(estimated_gain_factors),
                        places=6,
                    )
                    estimated_gain_factor = np.mean(estimated_gain_factors)
                    estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                        torch.tensor(estimated_gain_factor)
                    ).item()
                    est_gain_factors.append(estimated_gain_factor_in_db)

                    self.assertGreaterEqual(estimated_gain_factor_in_db, -30)
                    self.assertLessEqual(estimated_gain_factor_in_db, 0)

            if num_perturbed_channels == 0:
                perturbation_type_counter["zero_channels_changed"] += 1
            elif num_perturbed_channels == 1:
                if np.allclose(processed_samples[i, 0], samples_batch[i, 0]):
                    perturbation_type_counter["only_the_first_channel_changed"] += 1
                else:
                    perturbation_type_counter["only_the_second_channel_changed"] += 1
            elif num_perturbed_channels == 2:
                if np.allclose(est_gain_factors[0], est_gain_factors[1]):
                    # This should be very unlikely
                    perturbation_type_counter["both_channels_changed_equally"] += 1
                else:
                    perturbation_type_counter["both_channels_changed_differently"] += 1

        self.assertGreater(perturbation_type_counter["zero_channels_changed"], 100)
        self.assertLess(perturbation_type_counter["zero_channels_changed"], 500)
        self.assertGreater(
            perturbation_type_counter["only_the_first_channel_changed"], 100
        )
        self.assertLess(perturbation_type_counter["only_the_first_channel_changed"], 500)
        self.assertGreater(
            perturbation_type_counter["only_the_second_channel_changed"], 100
        )
        self.assertLess(perturbation_type_counter["only_the_second_channel_changed"], 500)
        self.assertLess(perturbation_type_counter["both_channels_changed_equally"], 10)
        self.assertGreater(
            perturbation_type_counter["both_channels_changed_differently"], 100
        )
        self.assertLess(
            perturbation_type_counter["both_channels_changed_differently"], 500
        )

        self.assertEqual(num_unprocessed_channels + num_processed_channels, 1000 * 2)
        self.assertGreater(num_processed_channels, 200 * 2)
        self.assertLess(num_processed_channels, 800 * 2)

        self.assertEqual(processed_samples.dtype, np.float32)

    def test_gain_per_channel_with_p_mode_per_batch(self):
        samples = np.array(
            [[1.0, 0.5, 0.25, 0.125, 0.01], [0.95, 0.5, 0.25, 0.12, 0.011]],
            dtype=np.float32,
        )
        samples_batch = np.stack([samples] * 100, axis=0)
        sample_rate = 16000

        augment = Gain(
            min_gain_in_db=-30.0,
            max_gain_in_db=-1.0,
            mode="per_channel",
            p=0.5,
            p_mode="per_batch",
        )
        num_processed_batches = 0
        for i in range(100):
            processed_samples = augment(
                samples=torch.from_numpy(samples_batch), sample_rate=sample_rate
            ).numpy()

            if np.allclose(processed_samples, samples_batch):
                continue
            else:
                num_processed_batches += 1

            num_unprocessed_channels = 0
            num_processed_channels = 0
            perturbation_type_counter = {
                "zero_channels_changed": 0,
                "only_the_first_channel_changed": 0,
                "only_the_second_channel_changed": 0,
                "both_channels_changed_differently": 0,
                "both_channels_changed_equally": 0,
            }
            for i in range(processed_samples.shape[0]):
                num_perturbed_channels = 0
                est_gain_factors = []
                for chn_index in range(processed_samples.shape[1]):
                    if np.allclose(
                        processed_samples[i, chn_index], samples_batch[i, chn_index]
                    ):
                        num_unprocessed_channels += 1
                    else:
                        num_processed_channels += 1
                        num_perturbed_channels += 1

                        estimated_gain_factors = (
                            processed_samples[i, chn_index] / samples_batch[i, chn_index]
                        )
                        self.assertAlmostEqual(
                            np.amin(estimated_gain_factors),
                            np.amax(estimated_gain_factors),
                            places=6,
                        )
                        estimated_gain_factor = np.mean(estimated_gain_factors)
                        estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                            torch.tensor(estimated_gain_factor)
                        ).item()
                        est_gain_factors.append(estimated_gain_factor_in_db)

                        self.assertGreaterEqual(estimated_gain_factor_in_db, -30)
                        self.assertLessEqual(estimated_gain_factor_in_db, 0)

                if num_perturbed_channels == 0:
                    perturbation_type_counter["zero_channels_changed"] += 1
                elif num_perturbed_channels == 1:
                    if np.allclose(processed_samples[i, 0], samples_batch[i, 0]):
                        perturbation_type_counter["only_the_first_channel_changed"] += 1
                    else:
                        perturbation_type_counter["only_the_second_channel_changed"] += 1
                elif num_perturbed_channels == 2:
                    if np.allclose(est_gain_factors[0], est_gain_factors[1]):
                        # This should be very unlikely
                        perturbation_type_counter["both_channels_changed_equally"] += 1
                    else:
                        perturbation_type_counter[
                            "both_channels_changed_differently"
                        ] += 1

            self.assertEqual(perturbation_type_counter["zero_channels_changed"], 0)
            self.assertLess(
                perturbation_type_counter["only_the_first_channel_changed"], 2
            )
            self.assertLess(
                perturbation_type_counter["only_the_second_channel_changed"], 2
            )
            self.assertLess(perturbation_type_counter["both_channels_changed_equally"], 2)
            self.assertGreater(
                perturbation_type_counter["both_channels_changed_differently"], 10
            )

            self.assertEqual(num_processed_channels, 100 * 2)
            self.assertEqual(processed_samples.dtype, np.float32)

        self.assertGreater(num_processed_batches, 10)
        self.assertLess(num_processed_batches, 90)

    def test_gain_per_channel_with_p_mode_per_example(self):
        samples = np.array(
            [[1.0, 0.5, 0.25, 0.125, 0.01], [0.95, 0.5, 0.25, 0.12, 0.011]],
            dtype=np.float32,
        )
        samples_batch = np.stack([samples] * 10000, axis=0)

        augment = Gain(
            min_gain_in_db=-30.0,
            max_gain_in_db=-1.0,
            mode="per_channel",
            p=0.5,
            p_mode="per_example",
        )
        processed_samples = augment(
            samples=torch.from_numpy(samples_batch), sample_rate=None
        ).numpy()

        num_unprocessed_examples = 0
        num_processed_examples = 0
        perturbation_type_counter = {
            "zero_channels_changed": 0,
            "only_the_first_channel_changed": 0,
            "only_the_second_channel_changed": 0,
            "both_channels_changed_differently": 0,
            "both_channels_changed_equally": 0,
        }
        for i in range(processed_samples.shape[0]):
            num_perturbed_channels = 0
            est_gain_factors = []

            if np.allclose(processed_samples[i], samples_batch[i]):
                num_unprocessed_examples += 1
            else:
                num_processed_examples += 1

            for chn_index in range(processed_samples.shape[1]):
                if not np.allclose(
                    processed_samples[i, chn_index], samples_batch[i, chn_index]
                ):
                    num_perturbed_channels += 1

                    estimated_gain_factors = (
                        processed_samples[i, chn_index] / samples_batch[i, chn_index]
                    )
                    self.assertAlmostEqual(
                        np.amin(estimated_gain_factors),
                        np.amax(estimated_gain_factors),
                        places=6,
                    )
                    estimated_gain_factor = np.mean(estimated_gain_factors)
                    estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                        torch.tensor(estimated_gain_factor)
                    ).item()
                    est_gain_factors.append(estimated_gain_factor_in_db)

                    self.assertGreaterEqual(estimated_gain_factor_in_db, -30)
                    self.assertLessEqual(estimated_gain_factor_in_db, 0)

            if num_perturbed_channels == 0:
                perturbation_type_counter["zero_channels_changed"] += 1
            elif num_perturbed_channels == 1:
                if np.allclose(processed_samples[i, 0], samples_batch[i, 0]):
                    perturbation_type_counter["only_the_first_channel_changed"] += 1
                else:
                    perturbation_type_counter["only_the_second_channel_changed"] += 1
            elif num_perturbed_channels == 2:
                if np.allclose(est_gain_factors[0], est_gain_factors[1]):
                    # This should be very unlikely
                    perturbation_type_counter["both_channels_changed_equally"] += 1
                else:
                    perturbation_type_counter["both_channels_changed_differently"] += 1

        self.assertGreater(perturbation_type_counter["zero_channels_changed"], 2000)
        self.assertLess(perturbation_type_counter["zero_channels_changed"], 8000)

        # We allow both these two to be 1 due to numerical accuracy. Ideally they should be 0.
        self.assertLessEqual(
            perturbation_type_counter["only_the_first_channel_changed"], 1
        )
        self.assertLessEqual(
            perturbation_type_counter["only_the_second_channel_changed"], 1
        )

        self.assertLess(perturbation_type_counter["both_channels_changed_equally"], 100)
        self.assertGreater(
            perturbation_type_counter["both_channels_changed_differently"], 2000
        )
        self.assertLess(
            perturbation_type_counter["both_channels_changed_differently"], 8000
        )

        self.assertEqual(num_unprocessed_examples + num_processed_examples, 10000)
        self.assertGreater(num_processed_examples, 2000)
        self.assertLess(num_processed_examples, 8000)

        self.assertEqual(processed_samples.dtype, np.float32)

    def test_gain_per_batch(self):
        samples = np.array(
            [[1.0, 0.5, 0.25, 0.125, 0.01], [0.95, 0.5, 0.25, 0.12, 0.011]],
            dtype=np.float32,
        )
        samples_batch = np.stack([samples] * 4, axis=0)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-30.0, max_gain_in_db=0.0, mode="per_batch", p=0.5)
        num_unprocessed_batches = 0
        num_processed_batches = 0
        for i in range(1000):
            processed_samples = augment(
                samples=torch.from_numpy(samples_batch), sample_rate=sample_rate
            ).numpy()

            estimated_gain_factors = processed_samples / samples_batch
            self.assertAlmostEqual(
                np.amin(estimated_gain_factors), np.amax(estimated_gain_factors), places=6
            )
            if np.allclose(processed_samples, samples_batch):
                num_unprocessed_batches += 1
            else:
                num_processed_batches += 1

        self.assertGreater(num_processed_batches, 200)
        self.assertGreater(num_unprocessed_batches, 200)

    def test_eval(self):
        samples = np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-15, max_gain_in_db=5, p=1.0)
        augment.eval()

        processed_samples = augment(
            samples=torch.from_numpy(samples), sample_rate=sample_rate
        ).numpy()

        np.testing.assert_array_equal(samples, processed_samples)

    def test_variability_within_batch(self):
        samples = np.array([[1.0, 0.5, 0.25, 0.125, 0.01]], dtype=np.float32)
        samples_batch = np.stack([samples] * 10000, axis=0)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-6, max_gain_in_db=6, p=0.5)
        processed_samples = augment(
            samples=torch.from_numpy(samples_batch), sample_rate=sample_rate
        ).numpy()
        self.assertEqual(processed_samples.dtype, np.float32)

        num_unprocessed_examples = 0
        num_processed_examples = 0
        actual_gains_in_db = []
        for i in range(processed_samples.shape[0]):
            if np.allclose(processed_samples[i], samples_batch[i]):
                num_unprocessed_examples += 1
            else:
                num_processed_examples += 1

                estimated_gain_factor = np.mean(processed_samples[i] / samples_batch[i])
                estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                    torch.tensor(estimated_gain_factor)
                ).item()

                self.assertGreaterEqual(estimated_gain_factor_in_db, -6)
                self.assertLessEqual(estimated_gain_factor_in_db, 6)
                actual_gains_in_db.append(estimated_gain_factor_in_db)

        mean_gain_in_db = np.mean(actual_gains_in_db)
        self.assertGreater(mean_gain_in_db, -1)
        self.assertLess(mean_gain_in_db, 1)

        self.assertEqual(num_unprocessed_examples + num_processed_examples, 10000)
        self.assertGreater(num_processed_examples, 2000)
        self.assertLess(num_processed_examples, 8000)

    def test_variability_within_batch_with_p_mode_per_batch(self):
        samples = np.array([[1.0, 0.5, 0.25, 0.125, 0.01]], dtype=np.float32)
        samples_batch = np.stack([samples] * 100, axis=0)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-6, max_gain_in_db=6, p=0.5, p_mode="per_batch")

        num_processed_batches = 0
        for _ in range(100):
            processed_samples = augment(
                samples=torch.from_numpy(samples_batch), sample_rate=sample_rate
            ).numpy()
            self.assertEqual(processed_samples.dtype, np.float32)

            if np.allclose(processed_samples, samples_batch):
                continue
            else:
                num_processed_batches += 1

            num_unprocessed_examples = 0
            num_processed_examples = 0
            actual_gains_in_db = []
            for i in range(processed_samples.shape[0]):
                if np.allclose(processed_samples[i], samples_batch[i]):
                    num_unprocessed_examples += 1
                else:
                    num_processed_examples += 1

                    estimated_gain_factor = np.mean(
                        processed_samples[i] / samples_batch[i]
                    )
                    estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                        torch.tensor(estimated_gain_factor)
                    ).item()

                    self.assertGreaterEqual(estimated_gain_factor_in_db, -6)
                    self.assertLessEqual(estimated_gain_factor_in_db, 6)
                    actual_gains_in_db.append(estimated_gain_factor_in_db)

            mean_gain_in_db = np.mean(actual_gains_in_db)
            self.assertGreater(mean_gain_in_db, -1.5)
            self.assertLess(mean_gain_in_db, 1.5)

            # Should be 0 and 100, but I give some slack due to possible numerical issues
            self.assertLessEqual(num_unprocessed_examples, 1)
            self.assertGreaterEqual(num_processed_examples, 99)

        self.assertGreater(num_processed_batches, 10)
        self.assertLess(num_processed_batches, 90)

    def test_reset_distribution(self):
        samples = np.array([[1.0, 0.5, 0.25, 0.125, 0.01]], dtype=np.float32)
        samples_batch = np.stack([samples] * 10000, axis=0)
        sample_rate = 16000

        augment = Gain(
            min_gain_in_db=-6, max_gain_in_db=6, p=0.5, sample_rate=sample_rate
        )
        # Change the parameters after init
        augment.min_gain_in_db = -18
        augment.max_gain_in_db = 3
        processed_samples = augment(samples=torch.from_numpy(samples_batch)).numpy()
        self.assertEqual(processed_samples.dtype, np.float32)

        actual_gains_in_db = []
        for i in range(processed_samples.shape[0]):
            if not np.allclose(processed_samples[i], samples_batch[i]):

                estimated_gain_factor = np.mean(processed_samples[i] / samples_batch[i])
                estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                    torch.tensor(estimated_gain_factor)
                ).item()

                self.assertGreaterEqual(estimated_gain_factor_in_db, -18)
                self.assertLessEqual(estimated_gain_factor_in_db, 3)
                actual_gains_in_db.append(estimated_gain_factor_in_db)

        mean_gain_in_db = np.mean(actual_gains_in_db)
        self.assertGreater(mean_gain_in_db, (-18 + 3) / 2 - 1)
        self.assertLess(mean_gain_in_db, (-18 + 3) / 2 + 1)

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
    def test_cuda_reset_distribution(self):
        samples = np.array([[1.0, 0.5, 0.25, 0.125, 0.01]], dtype=np.float32)
        samples_batch = np.stack([samples] * 10000, axis=0)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-6, max_gain_in_db=6, p=0.5).cuda()
        # Change the parameters after init
        augment.min_gain_in_db = -18
        augment.max_gain_in_db = 3
        processed_samples = (
            augment(
                samples=torch.from_numpy(samples_batch).cuda(), sample_rate=sample_rate
            )
            .cpu()
            .numpy()
        )
        self.assertEqual(processed_samples.dtype, np.float32)

        actual_gains_in_db = []
        for i in range(processed_samples.shape[0]):
            if not np.allclose(processed_samples[i], samples_batch[i]):

                estimated_gain_factor = np.mean(processed_samples[i] / samples_batch[i])
                estimated_gain_factor_in_db = convert_amplitude_ratio_to_decibels(
                    torch.tensor(estimated_gain_factor)
                ).item()

                self.assertGreaterEqual(estimated_gain_factor_in_db, -18)
                self.assertLessEqual(estimated_gain_factor_in_db, 3)
                actual_gains_in_db.append(estimated_gain_factor_in_db)

        mean_gain_in_db = np.mean(actual_gains_in_db)
        self.assertGreater(mean_gain_in_db, (-18 + 3) / 2 - 1)
        self.assertLess(mean_gain_in_db, (-18 + 3) / 2 + 1)

    def test_invalid_distribution(self):
        with self.assertRaises(ValueError):
            Gain(min_gain_in_db=18, max_gain_in_db=-3, p=0.5)

        augment = Gain(min_gain_in_db=-6, max_gain_in_db=-3, p=1.0)
        # Change the parameters after init
        augment.min_gain_in_db = 18
        augment.max_gain_in_db = 3
        with self.assertRaises(ValueError):
            augment(torch.tensor([[[1.0, 0.5, 0.25, 0.125]]], dtype=torch.float32), 16000)

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
    def test_gain_to_device_cuda(self):
        samples = np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32)
        sample_rate = 16000

        cuda_device = torch.device("cuda")

        augment = Gain(min_gain_in_db=-6.000001, max_gain_in_db=-6, p=1.0).to(
            device=cuda_device
        )
        processed_samples = (
            augment(
                samples=torch.from_numpy(samples).to(device=cuda_device),
                sample_rate=sample_rate,
            )
            .cpu()
            .numpy()
        )
        expected_factor = convert_decibels_to_amplitude_ratio(-6)
        assert_almost_equal(
            processed_samples,
            expected_factor
            * np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32),
            decimal=6,
        )
        self.assertEqual(processed_samples.dtype, np.float32)

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
    def test_gain_cuda(self):
        samples = np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-6.000001, max_gain_in_db=-6, p=1.0).cuda()
        processed_samples = (
            augment(samples=torch.from_numpy(samples).cuda(), sample_rate=sample_rate)
            .cpu()
            .numpy()
        )
        expected_factor = convert_decibels_to_amplitude_ratio(-6)
        assert_almost_equal(
            processed_samples,
            expected_factor
            * np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32),
            decimal=6,
        )
        self.assertEqual(processed_samples.dtype, np.float32)

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
    def test_gain_cuda_cpu(self):
        samples = np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32)
        sample_rate = 16000

        augment = Gain(min_gain_in_db=-6.000001, max_gain_in_db=-6, p=1.0).cuda().cpu()
        processed_samples = (
            augment(samples=torch.from_numpy(samples).cpu(), sample_rate=sample_rate)
            .cpu()
            .numpy()
        )
        expected_factor = convert_decibels_to_amplitude_ratio(-6)
        assert_almost_equal(
            processed_samples,
            expected_factor
            * np.array([[[1.0, 0.5, -0.25, -0.125, 0.0]]], dtype=np.float32),
            decimal=6,
        )
        self.assertEqual(processed_samples.dtype, np.float32)


import cv2
import json
import codecs
import os
import torch
import numpy as np
import time

import argparse

from config.ReID import cfg as ReIDCfg
from ReID_model.modeling import ReID_Model

from utils_BINGO.K_Means import k_means
from ReID_model.utils.dataset_loader import transform_imgs,ReID_imgs_loadv1

def imgs_sorted_by_ReID(ReID, ReIDCfg, imgs, distance_threshold=1, feat_norm='yes', version=0, input_features=None, input_gf=None,batch_size=50):
	# preprocess the imgs to ReID model form
	# imgs_array = transform_imgs(ReIDCfg,sub_imgs)

	with torch.no_grad():
		# if torch.cuda.device_count() >=1:
		# print('torch.cuda.device.count = {}'.format(torch.cuda.device_count()))
		if type(input_gf) == np.ndarray:
			feats = input_gf
			feats = torch.tensor(feats).cuda()
		else:
			imgs_length = len(imgs)
			leftover = 0
			if (imgs_length) % batch_size :
				leftover = 1
			num_batches = imgs_length // batch_size + leftover
			feats_ = []
			for j in range(num_batches):
				imgs_array_j = transform_imgs(ReIDCfg, imgs[j*batch_size:min((j+1)*batch_size , imgs_length)])
				imgs_array_j = imgs_array_j.to('cuda')
				feats_j = ReID(imgs_array_j)
				feats_.append(feats_j)
			feats = torch.cat(feats_, dim=0)
		if feat_norm == 'yes':
			# print("The test feature is normalized")
			feats = torch.nn.functional.normalize(feats, dim=1, p=2)
		# query : the main sub img is the first one!
		if version == 0:
			qf = feats[:1]
		elif version == 1:
			qf = torch.mean(feats, dim=0)
			qf = torch.unsqueeze(qf,dim=0)
		elif version == 2:
			qf = input_features
		# q_pids : np.asarray(self.pids[:self.num_query])
		# q_camids = np.asarray(self.camids[:self.num_query])
		# gallery : means the other sub imgs

		gf = feats[:]
		# g_pids = np.asarray(self.pids[self.num_query:])
		# g_camids = np.asarray(self.camids[self.num_query:])
		m, n = qf.shape[0], gf.shape[0]
		# print('m = {}, n = {}'.format(m, n))
		# calculate the distance between the query and gallery.
		# distance = x^2 + y^2 -2*x*y
		distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
		          torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
		distmat.addmm_(1, -2, qf, gf.t())
		distmat = distmat.cpu().numpy()
		# imgs_index is a tuple
		imgs_index = np.where(distmat < distance_threshold)

	return imgs_index[1].tolist(), distmat, qf

if __name__ == '__main__':
	os.environ["CUDA_VISIBLE_DEVICES"] = '1,5'
	import shutil

	parser = argparse.ArgumentParser(description='detection and tracking!')
	# set for ReID classification
	parser.add_argument(
		"--ReIDCfg",
		default="/datanew/hwb/Re-ID_models/reid-strong-baseline-master/configs/softmax_triplet_with_center.yml",
		help="path to config file", type=str)
	args = parser.parse_args()

	####################################
	# Load in ReID Classification model#
	####################################
	print('===> Start to constructing and loading ReID model', ['yellow', 'bold'])
	if args.ReIDCfg != "":
		ReIDCfg.merge_from_file(args.ReIDCfg)
	ReIDCfg.freeze()

	# imgs_arrays, img_names = ReID_imgs_load(ReIDCfg,main_img_dir)

	ReID = ReID_Model(ReIDCfg)
	ReID.cuda()
	# ReID = torch.nn.DataParallel(ReID)
	print('===> Finished constructing and loading ReID model', ['blue', 'bold'])

	# json_file = '/datanew/hwb/data/Football/SoftWare/9/with_VOT_test_1.json'
	# with codecs.open(json_file, 'r', 'utf-8-sig') as f:
	# 	action_datas = json.load(f)
	# 需要分类的
	# json_file_after_cluster = '/datanew/hwb/data/Football/SoftWare/0/with_VOT_test_0.json'
	game_id = 0
	input_gf = False
	wrong_index = [5 for i in range(100)]

	vis_dir_main = '/datanew/hwb/data/Football/SoftWare/{}/intermediate_results/vis_ReID/'.format(game_id)
	if os.path.isdir(vis_dir_main):
		shutil.rmtree(vis_dir_main)

	for dir_index in wrong_index:
		print(dir_index)
		main_img_dir = '/datanew/hwb/data/Football/SoftWare/{}/intermediate_results/Calibrate_transfer/{}/tracking'.format(game_id,dir_index)
		feature_path = '/datanew/hwb/data/Football/SoftWare/{}/intermediate_results/Calibrate_transfer/{}/tracking/{}_ReID_features_tracking.npy'.format(game_id,dir_index,dir_index)
		feature = np.load(feature_path)
		# feature = torch.tensor(feature).cuda()
		main_names = [name for name in os.listdir(main_img_dir) if name.split('.')[-1]=='jpg']
		main_names = sorted(main_names, key=lambda x:int(x.split('.')[0]))
		imgs = [cv2.imread(os.path.join(main_img_dir,img_name)) for img_name in main_names]

		main_img_dir_detection = '/datanew/hwb/data/Football/SoftWare/{}/intermediate_results/Calibrate_transfer/{}/detection'.format(game_id,dir_index)
		feature_path_detection = '/datanew/hwb/data/Football/SoftWare/{}/intermediate_results/Calibrate_transfer/{}/detection/{}_ReID_features_detection.npy'.format(game_id, dir_index,dir_index)
		feature_detection = np.load(feature_path_detection)
		# feature_detection = torch.tensor(feature_detection).cuda()

		main_names_detection = [name for name in os.listdir(main_img_dir_detection) if name.split('.')[-1] == 'jpg']
		main_names_detection = sorted(main_names_detection, key=lambda x: int(x.split('.')[0]))
		imgs_detection = [cv2.imread(os.path.join(main_img_dir_detection, img_name)) for img_name in main_names_detection]

		vis_dir = '/datanew/hwb/data/Football/SoftWare/{}/intermediate_results/vis_ReID/{}'.format(game_id,dir_index)
		if input_gf == True:
			imgs_index, distmat, output_feature = imgs_sorted_by_ReID(ReID, ReIDCfg, imgs, distance_threshold=1, feat_norm='yes', version=0, input_gf=feature)
		else:
			imgs_index, distmat, output_feature = imgs_sorted_by_ReID(ReID, ReIDCfg, imgs, distance_threshold=1, feat_norm='yes', version=0, input_gf=None)
		Positive_dir = os.path.join(vis_dir,'')
		os.makedirs(Positive_dir,exist_ok=True)
		Negative_dir = os.path.join(vis_dir,'Negative')
		os.makedirs(Negative_dir,exist_ok=True)

		for P_index, _ in enumerate(imgs):
			distance = distmat[0,P_index]
			if P_index in imgs_index:
				cv2.imwrite(os.path.join(Positive_dir,'{}_{:3f}.jpg'.format(P_index,distance)),imgs[P_index])
			else:
				cv2.imwrite(os.path.join(Negative_dir,'{}_{:3f}.jpg'.format(P_index,distance)),imgs[P_index])

		if input_gf == True:
			imgs_index_detection, distmat_detection, _ = imgs_sorted_by_ReID(ReID, ReIDCfg, imgs_detection, distance_threshold=1, feat_norm='yes',
														  version=2,input_features=output_feature,input_gf=feature_detection)
		else:
			imgs_index_detection, distmat_detection, _ = imgs_sorted_by_ReID(ReID, ReIDCfg, imgs_detection,
																			 distance_threshold=1, feat_norm='yes',
																			 version=2, input_features=output_feature,
																			 input_gf=None)

		Positive_dir_detection = os.path.join(vis_dir, 'detection')
		os.makedirs(Positive_dir_detection, exist_ok=True)
		Negative_dir_detection = os.path.join(vis_dir,'detection', 'Negative')
		os.makedirs(Negative_dir_detection, exist_ok=True)

		for P_index_detection, _ in enumerate(imgs_detection):
			distance = distmat_detection[0,P_index_detection]
			if P_index_detection in imgs_index_detection:
				cv2.imwrite(os.path.join(Positive_dir_detection, '{}_{:3f}.jpg'.format(P_index_detection,distance)), imgs_detection[P_index_detection])
			else:
				cv2.imwrite(os.path.join(Negative_dir_detection, '{}_{:3f}.jpg'.format(P_index_detection,distance)), imgs_detection[P_index_detection])


#coding:utf-8

"""
This is an experiment of PyAudio full duplex, record, filter, mix, and play simultaneously via ASIO.

Roland's old QUAD-CAPTURE has 4in/4out and can record and play simultaneously via ASIO.
Digtal audio music source is connected to COAXIAL IN (3-4) via SPDIF, FS=48KHz,
Analog signal like MIC is connected to  analog INPUT(1-2).
Via python scipy filter LPF4/HPF4, the music source mid sound is removed, and low and high sound remains.
And then, mix analog input(1-2) with the sound, and output from OUTPUT 1L/2R(1-2).
Turn MIX on QUAD-CAPTURE panel into PLAYBACK only.

music(Digital)---> COAXIAL IN(3-4) ---> Filter --|
                                                 |
MIC(Analog)   ---> INPUT 1L/2R(1-2) ----------> Mix --->  OUTPUT 1L/2R(1-2)

"""
"""
Copyright (c) 2020 Shun

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import sys
import time
import numpy as np
import pyaudio as pa
from scipy import signal # version > 1.2.0

# Check version
#  Python 3.6.4 on win32 (Windows 10)
#  numpy 1.18.4
#  scipy 1.4.1
#  PyAudio 0.2.11 from Unofficial Windows Binaries for Python Extension Packages Includes ASIO, DS, WMME, WASAPI, WDMKS support.

# global
fs48 = 48000
channel4 = 4
chunk1024 = 1024

class FILTER4(object):
    def __init__(self,fc1=150, fc2=5000, sr=48000 ):
        # fc1 low pass filter  cut-off frequency
        # fc2 high pass filter cut-off frequency
        # sr sampling rate
        self.sr=sr
        # compute filter coefficient of LPF4 and HPF4
        self.b_l4, self.a_l4 = signal.iirfilter(4,  fc1 ,  btype='lowpass',  ftype='butter', fs=self.sr )
        self.b_h4, self.a_h4 = signal.iirfilter(4,  fc2 ,  btype='highpass', ftype='butter', fs=self.sr )

        # initial value of LPF4 and HPF4
        self.zi_l4_l = np.zeros( (max(len(self.a_l4), len(self.b_l4)) - 1) )
        self.zi_l4_r = np.zeros( (max(len(self.a_l4), len(self.b_l4)) - 1) )
        self.zi_h4_l = np.zeros( (max(len(self.a_h4), len(self.b_h4)) - 1) )
        self.zi_h4_r = np.zeros( (max(len(self.a_h4), len(self.b_h4)) - 1) )

    def __call__(self, x_in):
        # turn to float
        in_float = np.frombuffer(x_in, dtype=np.int16).astype(np.float).reshape((chunk1024,channel4))

        # apply filter LPF4 HPF4 to coaxial in (3-4)
        # compute per one channel. If two channels at once, it's not work well.
        #
        y_l4_l, self.zi_l4_l= signal.lfilter(self.b_l4, self.a_l4, in_float[:,2], zi=self.zi_l4_l)
        y_l4_r, self.zi_l4_r= signal.lfilter(self.b_l4, self.a_l4, in_float[:,3], zi=self.zi_l4_r)
        y_h4_l, self.zi_h4_l= signal.lfilter(self.b_h4, self.a_h4, in_float[:,2], zi=self.zi_h4_l)
        y_h4_r, self.zi_h4_r= signal.lfilter(self.b_h4, self.a_h4, in_float[:,3], zi=self.zi_h4_r)

        # mix input(1-2) with above filter output, and output to OUTPUT(1,2)
        in_float[:,0] += y_l4_l + y_h4_l
        in_float[:,1] += y_l4_r + y_h4_r
        # If following, sometime soud drops, due to it maybe take over time to proceed.
        #
        #in_float[:,0] += y_l4_l + y_h4_l + in_float[:,0]
        #in_float[:,1] += y_l4_r + y_h4_r + in_float[:,1]

        # back to int16
        return bytes(np.array(in_float, dtype=np.int16))

# instance
flt_stream= FILTER4(sr=fs48)

# define callback function, PyAudio non-blocking
def callback_rec_play(in_data, frame_count, time_info, status):
    in_data= flt_stream( in_data)
    return (in_data, pa.paContinue)

def search_device( dname=None, host_name=None):
    # get usb index number of specfied spec.
    use_device_index=None

    # PyAudio
    p_in = pa.PyAudio()
    for i in range(p_in.get_device_count()):

        devinfo= p_in.get_device_info_by_index(i)

        # get host api's name
        for k in list(devinfo.items()):
            name, value = k
            if name == 'hostApi':
            	host_api= p_in.get_host_api_info_by_index(k[1])['name']

        if dname in devinfo['name'] and host_name in host_api:
            use_device_index= i
            print ('use_device_index', use_device_index)
            print ('')

    # PyAudio terminate
    p_in.terminate()
    return use_device_index

if __name__ == "__main__":

    # get usb device index
    use_device_index_inout1= search_device(dname='QUAD-CAPTURE',  host_name='ASIO')
    if use_device_index_inout1 is None:
        print ('ERROR: There is not QUAD-CAPTURE ASIO.')
        sys.exit()

    # PyAudio
    p_in1 = pa.PyAudio()
    py_format_inout1 = p_in1.get_format_from_width(2)  # 16bit

    # define stream as full duplex, non-blocking
    inout_stream1 = p_in1.open(format=py_format_inout1,
                      channels=channel4,
                      rate=fs48,
                      input=True,
                      output=True,
                      frames_per_buffer=chunk1024,
                      input_device_index=use_device_index_inout1,
                      output_device_index=use_device_index_inout1,
                      stream_callback=callback_rec_play)

    # start stream
    inout_stream1.start_stream()

    print ('stream starts..., Input any key and return to stop.')
    # input loop
    while inout_stream1.is_active():
        c = input()
        if c:
            break
        time.sleep(0.1)

    """
    This input loop is not loop.
    c = input() is waitting for return key.
    And time.sleep(0.1) does not called.
    Following is another sample code for Windows.

    import msvcrt  # for windows only

    # input loop
    while inout_stream1.is_active():
        if msvcrt.kbhit():   # if keyboard hit
            c=msvcrt.getch()
            break
        time.sleep(0.1)

    """

    # stream stop and close
    inout_stream1.stop_stream()
    inout_stream1.close()

    # PyAudio terminate
    p_in1.terminate()


"""PixelCNN on Flickr8k.
"""
"""
PixelCNN Training code and utilities are licensed under APL2.0 from

Parag Mital
---------------------
https://github.com/pkmital/pycadl/blob/master/cadl/pixelcnn.py

Copyright 2018 Holly Grimm.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import os
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from cadl import pixelcnn
from cadl import dataset_utils as dsu

def train(ckpt_path='pixelcnn',
                        n_epochs=1000,
                        save_step=100,
                        write_step=25,
                        B=4,
                        H=64,
                        W=64,
                        C=3):
    """Summary

    Parameters
    ----------
    ckpt_path : str, optional
        Description
    n_epochs : int, optional
        Description
    save_step : int, optional
        Description
    write_step : int, optional
        Description
    B : int, optional
        Description
    H : int, optional
        Description
    W : int, optional
        Description
    C : int, optional
        Description
    """
    ckpt_name = os.path.join(ckpt_path, 'pixelcnn.ckpt')

    with tf.Graph().as_default(), tf.Session() as sess:
        # Not actually conditioning on anything here just using the gated cnn model
        net = pixelcnn.build_conditional_pixel_cnn_model(B=B, H=H, W=W, C=C)

        # build the optimizer (this will take a while!)
        optimizer = tf.train.AdamOptimizer(
            learning_rate=0.001).minimize(net['cost'])

        # Load a list of files for flickr8k-64 dataset
        path = './flickr8k-64/'
        imagepath = os.path.join(path, 'Flicker8k_Dataset')
        tokenFilename =  os.path.join(path, 'Flickr8k.token.txt')

        with open(tokenFilename) as f:
            captiontxt = f.readlines()
        fs = []
        labels = []

        captiontxt_everyfifthline = captiontxt[::5]

        for line in captiontxt_everyfifthline:
            parts = line.split('\t')
            imgid = parts[0].split('.jpg')[0]
            label = parts[1][:-1].replace('\n','').strip()
            fs.append(os.path.join(imagepath, imgid + '.jpg'))
            labels.append(label)

        # Create a threaded image pipeline which will load/shuffle/crop/resize
        batch = dsu.create_input_pipeline(
            fs,
            batch_size=B,
            n_epochs=n_epochs,
            shape=[64, 64, 3],
            crop_shape=[H, W, C],
            crop_factor=1.0,
            n_threads=8)

        saver = tf.train.Saver()
        writer = tf.summary.FileWriter(ckpt_path)
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())
        sess.run(init_op)

        # This will handle our threaded image pipeline
        coord = tf.train.Coordinator()

        # Ensure no more changes to graph
        tf.get_default_graph().finalize()

        # Start up the queues for handling the image pipeline
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        if os.path.exists(ckpt_name + '.index') or os.path.exists(ckpt_name):
            saver.restore(sess, ckpt_name)
            saver.restore(sess, tf.train.latest_checkpoint(ckpt_path))

        epoch_i = 0
        batch_i = 0
        try:
            while not coord.should_stop() and epoch_i < n_epochs:
                batch_i += 1
                batch_xs = sess.run(batch)
                train_cost = sess.run(
                    [net['cost'], optimizer], feed_dict={net['X']: batch_xs})[0]

                print(batch_i, train_cost)
                if batch_i % write_step == 0:
                    summary = sess.run(
                        net['summaries'], feed_dict={net['X']: batch_xs})
                    writer.add_summary(summary, batch_i)

                if batch_i % save_step == 0:
                    # Save the variables to disk.  Don't write the meta graph
                    # since we can use the code to create it, and it takes a long
                    # time to create the graph since it is so deep
                    saver.save(
                        sess,
                        ckpt_name,
                        global_step=batch_i,
                        write_meta_graph=True)
        except tf.errors.OutOfRangeError:
            print('Done.')
        finally:
            # One of the threads has issued an exception.  So let's tell all the
            # threads to shutdown.
            coord.request_stop()

        # Wait until all threads have finished.
        coord.join(threads)

def generate(img_index=1):
    """Use PixelCNN to regenerate the bottom half of a selected image
    Parameters
    ----------
    img_index : int, optional
        Index into the list of image files
    """
    # Parameters for generation
    ckpt_path = 'pixelcnn'
    B = None
    H = 64
    W = 64
    C = 3

    with tf.Graph().as_default(), tf.Session() as sess:
        # Not actually conditioning on anything here just using the gated cnn model
        net = pixelcnn.build_conditional_pixel_cnn_model(B=B, H=H, W=W, C=C)

        # Load a list of files for flickr8k
        path = './flickr8k-64/'
        imagepath = os.path.join(path, 'Flicker8k_Dataset')
        tokenFilename =  os.path.join(path, 'Flickr8k.token.txt')

        with open(tokenFilename) as f:
            captiontxt = f.readlines()
        fs = []
        labels = []

        captiontxt_everyfifthline = captiontxt[::5]

        for line in captiontxt_everyfifthline:
            parts = line.split('\t')
            imgid = parts[0].split('.jpg')[0]
            label = parts[1][:-1].replace('\n','').strip()
            fs.append(os.path.join(imagepath, imgid + '.jpg'))
            labels.append(label)

        saver = tf.train.Saver()
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())
        sess.run(init_op)
        saver.restore(sess, tf.train.latest_checkpoint(ckpt_path))

        # select one file to synthesize
        img = plt.imread(fs[img_index])
        from scipy.misc import imresize
        og_img = imresize(img, (H, W))
        img = og_img.copy()
        # Zero out bottom half of image and let's try to synthesize it
        img[H // 2:, :, :] = 0
        for h_i in range(H // 2, H):
            for w_i in range(W):

#         # Zero out right half of image and let's try to synthesize it
#         img[:, W // 2:, :] = 0
#         for h_i in range(H):
#             for w_i in range(W // 2, W):
                for c_i in range(C):
                    print(h_i, w_i, c_i, end='\r')
                    X = img.copy()
                    preds = sess.run(
                        net['sampled_preds'],
                        feed_dict={net['X']: X[np.newaxis]})
                    X = preds.reshape((1, H, W, C)).astype(np.uint8)
                    img[h_i, w_i, c_i] = X[0, h_i, w_i, c_i]

        fig, axs = plt.subplots(1, 2)
        axs[0].imshow(og_img)
        axs[1].imshow(img)
        plt.show()

if __name__ == '__main__':
    train()


"""Module to skim the PDB for similar unit cells"""

from __future__ import division

__author__ = "Adam Simpkin & Felix Simkovic"
__date__ = "30 Jun 2017"
__version__ = "1.0"

import ast
import datetime
import glob
import logging
import numpy as np
import os

from simbad.core.lattice_score import LatticeSearchResult
import simbad.util.pdb_util

logger = logging.getLogger(__name__)

class LatticeSearch(object):
    """A class to do a search for PDB entries with similar unit cell dimensions."""

    class _ResultCache(object):
        def __init__(self):
            self._codes = set()
            self._data = []

        def __iter__(self):
            for e in self._data:
                yield e

        def append(self, score):
            code = score.pdb_code.upper()
            if code in self._codes:
                return
            self._codes.add(code)
            self._data.append(score)

    def __init__(self, lattice_db_fname, model_dir):
        """Initialize a new Lattice Search class

        Parameters
        ----------
        lattice_db_fname : str
           The path to the lattice database [pickle format]
        max_to_keep : int, optional
           The maximum number of results to keep [default: 20]

        """
        self._lattice_db_fname = None

        self.lattice_db_fname = lattice_db_fname
        self.model_dir = model_dir

        self.results = []

    @property
    def lattice_db_fname(self):
        """The path to the lattice database"""
        return self._lattice_db_fname

    @lattice_db_fname.setter
    def lattice_db_fname(self, lattice_db_fname):
        """Define the lattice database filename"""
        # Check how old the database is, if older than 90 days print
        # message to suggest updating
        timestamp = os.path.getmtime(lattice_db_fname)
        if (datetime.date.today() - datetime.date.fromtimestamp(timestamp)).days > 90:
            logger.info("Lattice database is older than 90 days, consider updating!\n" 'Use the command "simbad-database lattice" in your Terminal')
        self._lattice_db_fname = lattice_db_fname

    def search(self, space_group, unit_cell, tolerance=0.05, max_to_keep=50, max_penalty=12):
        """Search for similar Niggli cells

        Parameters
        ----------
        unit_cell : list, tuple
           The parameters of the unit cell
        space_group : str
           The space group
        tolerance : int, float, optional
           The tolerance applied for Niggli cell comparison [default: 0.05]
        max_to_keep : int, optional
           The top-N number of results to return [default: 50]
        max_penalty : int, optional
           The total penalty score over which results are ignored [default: 12]

        """
        space_group = self.check_sg(space_group)
        niggli_cell = self.calculate_niggli_cell(unit_cell, space_group)
        niggli_cell = np.array(niggli_cell)

        tol_niggli_cell = niggli_cell * tolerance

        results = self._ResultCache()
        with np.load(self.lattice_db_fname) as compressed:
            for entry in compressed["arr_0"]:
                pdb_code = "".join(chr(c) for c in entry[:4].astype("uint8"))
                pdb_path = os.path.join(self.model_dir, "{}.pdb".format(pdb_code))
                alt_cell = chr(int(entry[4])) if entry[4] != 0.0 else " "
                db_cell = entry[5:]

                if self.cell_within_tolerance(niggli_cell, db_cell, tol_niggli_cell):
                    total_pen, length_pen, angle_pen = self.calculate_penalties(niggli_cell, db_cell)
                    vol_diff = self.calculate_volume_difference(niggli_cell, db_cell)
                    if total_pen < max_penalty:
                        prob = self.calculate_probability(total_pen)
                        score = LatticeSearchResult(pdb_code, pdb_path, alt_cell, db_cell, vol_diff, total_pen, length_pen, angle_pen, prob)
                        results.append(score)

        results_sorted = sorted(results, key=lambda x: float(x.total_penalty), reverse=False)
        self.results = results_sorted[:max_to_keep]

    @classmethod
    def calculate_penalties(cls, query, reference):
        """Calculate the linear cell variation between unit cells

        Parameters
        ----------
        query : list, tuple
           The query cell parameters
        reference : list, tuple
           The reference cell parameters

        Returns
        -------
        float
           Total penalty
        float
           Length penalty
        float
           Angle penalty
        """
        delta = np.absolute(query - reference)
        length_penalty = delta[:3].sum().item()
        angle_penalty = delta[3:].sum().item()
        total_penalty = length_penalty + angle_penalty
        return total_penalty, length_penalty, angle_penalty

    @classmethod
    def calculate_probability(cls, penalty_score):
        """Calculate the probability that a penalty score will give a solution

        Parameters
        ----------
        penalty_score : float
            The total penalty score calculate for a search result

        Returns
        -------
        float
            Probability score
        """
        x = -1.01 * penalty_score + 2.11
        return np.around(1 / (1 + np.exp(-x)), decimals=3)

    @classmethod
    def cell_within_tolerance(cls, query, reference, tolerance):
        """Compare two cells and determine if ``query`` is within ``reference`` cell parameter tolerance

        Parameters
        ----------
        query : list, tuple
           The query cell parameters
        reference : list, tuple
           The reference cell parameters
        tolerance : list, tuple
           The tolerance cell parameter values

        Returns
        -------
        bool

        """
        return np.all(np.absolute(query - reference) <= tolerance)

    @classmethod
    def calculate_volume_difference(cls, query, reference):
        """Calculate the difference in volume between the query unit cell and the reference unit cell

        Parameters
        ----------
        query : list, tuple
           The query cell parameters
        reference : list, tuple
           The reference cell parameters

        Returns
        -------
        float
            The absolute difference in cell volumes
        """
        import cctbx.uctbx
        cell_volume_1 = cctbx.uctbx.unit_cell(list(query)).volume()
        cell_volume_2 = cctbx.uctbx.unit_cell(list(reference)).volume()
        return np.absolute(cell_volume_1 - cell_volume_2).round(decimals=3).item()

    @staticmethod
    def calculate_niggli_cell(unit_cell, space_group):
        """Calculate the parameters of the Niggli cell

        Parameters
        ----------
        unit_cell : list, tuple
           The parameters of the unit cell
        space_group : str
           The space group

        Returns
        -------
        list
           The Niggli cell parameters

        """
        import cctbx.crystal
        import cctbx.uctbx

        unit_cell = cctbx.uctbx.unit_cell(list(unit_cell))
        xs = cctbx.crystal.symmetry(unit_cell=unit_cell, space_group=space_group, correct_rhombohedral_setting_if_necessary=True)
        niggli_cell = xs.change_basis(xs.change_of_basis_op_to_niggli_cell()).unit_cell()
        niggli_cell = list(ast.literal_eval(str(niggli_cell)))
        logger.info("Niggli cell calculated as: [%s]", ", ".join(map(str, niggli_cell)))
        return niggli_cell

    @staticmethod
    def check_sg(sg):
        """Check the space group for known anomalies"""
        sg_conversion = {
            "A1": "P1",
            "B2": "B112",
            "C1211": "C2",
            "F422": "I422",
            "I21": "I2",
            "I1211": "I2",
            "P21212A": "P212121",
            "R3": "R3:R",
            "C4212": "P422",
        }
        return sg_conversion.get(sg, sg)

    def copy_results(self, source, destination):
        """Copy the results from a local copy of the PDB

        Parameters
        ----------
        source : str
           The path to copy results from
        destination : str
           The path to save results to

        Raises
        ------
        ValueError
           No search results found/available
        ValueError
           Output directory does not exist
        IOError
            Search result not found in installed PDB

        """
        if not self.results:
            raise ValueError("No search results found/available")

        if not os.path.isdir(destination):
            msg = "Output directory does not exist: {0}".format(destination)
            raise ValueError(msg)

        prefix, ext = os.path.basename(glob.glob(os.path.join(source, "*", "*"))[0]).split(".", 1)
        prefix = prefix[:-4]

        to_del = []
        for count, result in enumerate(self.results):
            f_name = os.path.join(source, "{0}", "{1}{2}.{3}").format(result.pdb_code[1:3].lower(), prefix, result.pdb_code.lower(), ext)
            f_name_out = os.path.join(destination, "{0}.pdb".format(result.pdb_code))
            try:
                struct = simbad.util.pdb_util.PdbStructure.from_file(f_name)
                struct.standardize()
                struct.save(f_name_out)
            except IOError:
                logger.warning("Encountered problem copying PDB %s from %s - removing entry from list", result.pdb_code, source)
                to_del.append(count)
            except RuntimeError:
                logger.warning("Encountered problem copying PDB %s from %s - removing entry from list", result.pdb_code, source)
                to_del.append(count)

        for i in reversed(to_del):
            self.results.pop(i)

    def download_results(self, destination):
        """Download the results directly from the PDB

        Parameters
        ----------
        destination : str
           The path to save results to

        Raises
        ------
        ValueError
           No search results found/available
        ValueError
           Output directory does not exist
        RuntimeError
            Unable to download PDB

        """
        if not self.results:
            raise ValueError("No search results found/available")

        if not os.path.isdir(destination):
            msg = "Output directory does not exist: {0}".format(destination)
            raise ValueError(msg)

        to_del = []
        for count, result in enumerate(self.results):
            try:
                f_name_out = os.path.join(destination, result.pdb_code + ".pdb")
                struct = simbad.util.pdb_util.PdbStructure.from_pdb_code(result.pdb_code)
                struct.standardize()
                struct.save(f_name_out)
            except RuntimeError:
                logger.warning("Encountered problem downloading PDB %s - removing entry from list", result.pdb_code)
                to_del.append(count)

        for i in reversed(to_del):
            self.results.pop(i)

    def summarize(self, csvfile):
        """Summarize the search results

        Parameters
        ----------
        csvfile : str
           The path to an output CSV file

        """
        from simbad.util import summarize_result

        columns = [
            "alt",
            "a",
            "b",
            "c",
            "alpha",
            "beta",
            "gamma",
            "length_penalty",
            "angle_penalty",
            "total_penalty",
            "volume_difference",
            "probability_score",
        ]
        summarize_result(self.results, csv_file=csvfile, columns=columns)


import numpy as np
import scipy.io as scio
import os
import argparse
import pickle
from sklearn import metrics
import json
import socket

# data folder contain all datasets, such as ped1, ped2, avenue, shanghaitech, etc
DATA_DIR = '../Data'
# hostname = socket.gethostname()
# if hostname == 'dl-T8520-G10':  # 119
#     DATA_DIR = '/home/liuwen/ssd/datasets'
# elif hostname == 'admin' or hostname == 'compute101' or hostname == 'compute113' or hostname == 'compute106' \
#         or hostname == 'compute107' or hostname == 'compute114':   # node02
#     DATA_DIR = '/home/luowx/liuwen/datasets'
# elif hostname == 'gpu13' or 'gpu14':
#     DATA_DIR = '/public/home/gaoshenghua/liuwen/datasets'
# else:
#     # raise NotImplementedError('Not found this machine {}!'.format(hostname))
#     DATA_DIR = '../Data'

# normalize scores in each sub video
NORMALIZE = True

# number of history frames, since in prediction based method, the first 4 frames can not be predicted, so that
# the first 4frames are undecidable, we just ignore the first 4 frames
DECIDABLE_IDX = 4

def parser_args():
    parser = argparse.ArgumentParser(description='evaluating the model, computing the roc/auc.')

    parser.add_argument('-f', '--file', type=str, help='the path of loss file.')
    parser.add_argument('-t', '--type', type=str, default='compute_auc',
                        help='the type of evaluation, choosing type is: plot_roc, compute_auc, '
                             'test_func\n, the default type is compute_auc')
    return parser.parse_args()

class RecordResult(object):
    def __init__(self, fpr=None, tpr=None, auc=-np.inf, dataset=None, loss_file=None):
        self.fpr = fpr
        self.tpr = tpr
        self.auc = auc
        self.dataset = dataset
        self.loss_file = loss_file

    def __lt__(self, other):
        return self.auc < other.auc

    def __gt__(self, other):
        return self.auc > other.auc

    def __str__(self):
        return 'dataset = {}, loss file = {}, auc = {}'.format(self.dataset, self.loss_file, self.auc)

class GroundTruthLoader(object):
    AVENUE = 'avenue'
    PED1 = 'ped1'
    PED1_PIXEL_SUBSET = 'ped1_pixel_subset'
    PED2 = 'ped2'
    ENTRANCE = 'enter'
    EXIT = 'exit'
    SHANGHAITECH = 'shanghaitech'
    SHANGHAITECH_LABEL_PATH = os.path.join(DATA_DIR, 'shanghaitech/testing/test_frame_mask')
    TOY_DATA = 'toydata'
    TOY_DATA_LABEL_PATH = os.path.join(DATA_DIR, TOY_DATA, 'toydata.json')

    AIHUB = 'mini'
    AIHUB_LABEL_PATH = os.path.join(DATA_DIR,'mini/gt')

    NAME_MAT_MAPPING = {
        AVENUE: os.path.join(DATA_DIR, 'avenue/avenue.mat'),
        PED1: os.path.join(DATA_DIR, 'ped1/ped1.mat'),
        PED2: os.path.join(DATA_DIR, 'ped2/ped2.mat'),
        ENTRANCE: os.path.join(DATA_DIR, 'enter/enter.mat'),
        EXIT: os.path.join(DATA_DIR, 'exit/exit.mat')
    }

    NAME_FRAMES_MAPPING = {
        AVENUE: os.path.join(DATA_DIR, 'avenue/testing/frames'),
        PED1: os.path.join(DATA_DIR, 'ped1/testing/frames'),
        PED2: os.path.join(DATA_DIR, 'ped2/testing/frames'),
        ENTRANCE: os.path.join(DATA_DIR, 'enter/testing/frames'),
        EXIT: os.path.join(DATA_DIR, 'exit/testing/frames')
    }

    def __init__(self, mapping_json=None):
        """
        Initial a ground truth loader, which loads the ground truth with given dataset name.

        :param mapping_json: the mapping from dataset name to the path of ground truth.
        """

        if mapping_json is not None:
            with open(mapping_json, 'rb') as json_file:
                self.mapping = json.load(json_file)
        else:
            self.mapping = GroundTruthLoader.NAME_MAT_MAPPING

    def __call__(self, dataset):
        """ get the ground truth by provided the name of dataset.

        :type dataset: str
        :param dataset: the name of dataset.
        :return: np.ndarray, shape(#video)
                 np.array[0] contains all the start frame and end frame of abnormal events of video 0,
                 and its shape is (#frapsnr, )
        """

        if dataset == GroundTruthLoader.SHANGHAITECH:
            gt = self.__load_shanghaitech_gt()
        elif dataset == GroundTruthLoader.TOY_DATA:
            gt = self.__load_toydata_gt()
        elif dataset == GroundTruthLoader.AIHUB:
            gt = self.__load_aihub_gt()
        else:
            gt = self.__load_ucsd_avenue_subway_gt(dataset)
        return gt

    def __load_ucsd_avenue_subway_gt(self, dataset):
        assert dataset in self.mapping, 'there is no dataset named {} \n Please check {}' \
            .format(dataset, GroundTruthLoader.NAME_MAT_MAPPING.keys())

        mat_file = self.mapping[dataset]
        abnormal_events = scio.loadmat(mat_file, squeeze_me=True)['gt']

        if abnormal_events.ndim == 2:
            abnormal_events = abnormal_events.reshape(-1, abnormal_events.shape[0], abnormal_events.shape[1])

        num_video = abnormal_events.shape[0]
        dataset_video_folder = GroundTruthLoader.NAME_FRAMES_MAPPING[dataset]
        video_list = os.listdir(dataset_video_folder)
        video_list.sort()

        assert num_video == len(video_list), 'ground true does not match the number of testing videos. {} != {}' \
            .format(num_video, len(video_list))

        # get the total frames of sub video
        def get_video_length(sub_video_number):
            # video_name = video_name_template.format(sub_video_number)
            video_name = os.path.join(dataset_video_folder, video_list[sub_video_number])
            assert os.path.isdir(video_name), '{} is not directory!'.format(video_name)

            length = len(os.listdir(video_name))

            return length

        # need to test [].append, or np.array().append(), which one is faster
        gt = []
        for i in range(num_video):
            length = get_video_length(i)

            sub_video_gt = np.zeros((length,), dtype=np.int8)
            sub_abnormal_events = abnormal_events[i]
            if sub_abnormal_events.ndim == 1:
                sub_abnormal_events = sub_abnormal_events.reshape((sub_abnormal_events.shape[0], -1))

            _, num_abnormal = sub_abnormal_events.shape

            for j in range(num_abnormal):
                # (start - 1, end - 1)
                start = sub_abnormal_events[0, j] - 1
                end = sub_abnormal_events[1, j]

                sub_video_gt[start: end] = 1

            gt.append(sub_video_gt)

        return gt

    @staticmethod
    def __load_shanghaitech_gt():
        video_path_list = os.listdir(GroundTruthLoader.SHANGHAITECH_LABEL_PATH)
        video_path_list.sort()

        gt = []
        for video in video_path_list:
            # print(os.path.join(GroundTruthLoader.SHANGHAITECH_LABEL_PATH, video))
            gt.append(np.load(os.path.join(GroundTruthLoader.SHANGHAITECH_LABEL_PATH, video)))

        return gt

    @staticmethod
    def __load_aihub_gt():
        video_path_list = os.listdir(GroundTruthLoader.AIHUB_LABEL_PATH)

        video_path_list.sort()
        gt = []

        for video in video_path_list:
            gt.append(np.load(os.path.join(GroundTruthLoader.AIHUB_LABEL_PATH, video)))

        return gt

    @staticmethod
    def __load_toydata_gt():

        with open(GroundTruthLoader.TOY_DATA_LABEL_PATH, 'r') as gt_file:
            gt_dict = json.load(gt_file)

        gt = []
        for video, video_info in gt_dict.items():
            length = video_info['length']
            video_gt = np.zeros((length,), dtype=np.uint8)
            sub_gt = np.array(np.matrix(video_info['gt']))

            for anomaly in sub_gt:
                start = anomaly[0]
                end = anomaly[1] + 1
                video_gt[start: end] = 1
            gt.append(video_gt)
        return gt

    @staticmethod
    def get_pixel_masks_file_list(dataset):
        # pixel mask folder
        pixel_mask_folder = os.path.join(DATA_DIR, dataset, 'pixel_masks')
        pixel_mask_file_list = os.listdir(pixel_mask_folder)
        pixel_mask_file_list.sort()

        # get all testing videos
        dataset_video_folder = GroundTruthLoader.NAME_FRAMES_MAPPING[dataset]
        video_list = os.listdir(dataset_video_folder)
        video_list.sort()

        # get all testing video names with pixel masks
        pixel_video_ids = []
        ids = 0
        for pixel_mask_name in pixel_mask_file_list:
            while ids < len(video_list):
                if video_list[ids] + '.npy' == pixel_mask_name:
                    pixel_video_ids.append(ids)
                    ids += 1
                    break
                else:
                    ids += 1

        assert len(pixel_video_ids) == len(pixel_mask_file_list)

        for i in range(len(pixel_mask_file_list)):
            pixel_mask_file_list[i] = os.path.join(pixel_mask_folder, pixel_mask_file_list[i])

        return pixel_mask_file_list, pixel_video_ids

def load_psnr_gt(loss_file):
    with open(loss_file, 'rb') as reader:
        # results {
        #   'dataset': the name of dataset
        #   'psnr': the psnr of each testing videos,
        # }

        # psnr_records['psnr'] is np.array, shape(#videos)
        # psnr_records[0] is np.array   ------>     01.avi
        # psnr_records[1] is np.array   ------>     02.avi
        #               ......
        # psnr_records[n] is np.array   ------>     xx.avi

        results = pickle.load(reader)

    dataset = results['dataset']
    psnr_records = results['psnr']

    num_videos = len(psnr_records)

    # load ground truth
    gt_loader = GroundTruthLoader()
    gt = gt_loader(dataset=dataset)

    assert num_videos == len(gt), 'the number of saved videos does not match the ground truth, {} != {}' \
        .format(num_videos, len(gt))

    return dataset, psnr_records, gt

def load_psnr_gt_flow(loss_file):
    with open(loss_file, 'rb') as reader:
        # results {
        #   'dataset': the name of dataset
        #   'psnr': the psnr of each testing videos,
        # }

        # psnr_records['psnr'] is np.array, shape(#videos)
        # psnr_records[0] is np.array   ------>     01.avi
        # psnr_records[1] is np.array   ------>     02.avi
        #               ......
        # psnr_records[n] is np.array   ------>     xx.avi

        results = pickle.load(reader)

    dataset = results['dataset']
    psnrs = results['psnr']
    flows = results['flow']

    num_videos = len(psnrs)

    # load ground truth
    gt_loader = GroundTruthLoader()
    gt = gt_loader(dataset=dataset)

    assert num_videos == len(gt), 'the number of saved videos does not match the ground truth, {} != {}' \
        .format(num_videos, len(gt))

    return dataset, psnrs, flows, gt

def load_psnr(loss_file):
    """
    load image psnr or optical flow psnr.
    :param loss_file: loss file path
    :return:
    """
    with open(loss_file, 'rb') as reader:
        # results {
        #   'dataset': the name of dataset
        #   'psnr': the psnr of each testing videos,
        # }

        # psnr_records['psnr'] is np.array, shape(#videos)
        # psnr_records[0] is np.array   ------>     01.avi
        # psnr_records[1] is np.array   ------>     02.avi
        #               ......
        # psnr_records[n] is np.array   ------>     xx.avi

        results = pickle.load(reader)
    psnrs = results['psnr']
    return psnrs

def get_scores_labels(loss_file):
    # the name of dataset, loss, and ground truth
    dataset, psnr_records, gt = load_psnr_gt(loss_file=loss_file)

    # the number of videos
    num_videos = len(psnr_records)

    scores = np.array([], dtype=np.float32)
    labels = np.array([], dtype=np.int8)
    # video normalization
    for i in range(num_videos):
        distance = psnr_records[i]

        if NORMALIZE:
            distance -= distance.min()  # distances = (distance - min) / (max - min)
            distance /= distance.max()
            # distance = 1 - distance

        scores = np.concatenate((scores[:], distance[DECIDABLE_IDX:]), axis=0)
        labels = np.concatenate((labels[:], gt[i][DECIDABLE_IDX:]), axis=0)
    return dataset, scores, labels

def precision_recall_auc(loss_file):
    if not os.path.isdir(loss_file):
        loss_file_list = [loss_file]
    else:
        loss_file_list = os.listdir(loss_file)
        loss_file_list = [os.path.join(loss_file, sub_loss_file) for sub_loss_file in loss_file_list]

    optimal_results = RecordResult()
    for sub_loss_file in loss_file_list:
        dataset, scores, labels = get_scores_labels(sub_loss_file)
        precision, recall, thresholds = metrics.precision_recall_curve(labels, scores, pos_label=0)
        auc = metrics.auc(recall, precision)

        result_path = os.getcwd()+'/result_auc/'
        if not os.path.exists(result_path):
            os.makedirs(result_path)
        auprc ={}
        auprc['recall'] = recall
        auprc['precision'] = precision
        scio.savemat(result_path+'future_pred_auprc.mat',{'auprc':auprc})

        results = RecordResult(recall, precision, auc, dataset, sub_loss_file)

        if optimal_results < results:
            optimal_results = results

        if os.path.isdir(loss_file):
            print(results)
    print('##### optimal result and model = {}'.format(optimal_results))
    return optimal_results

def cal_eer(fpr, tpr):
    # makes fpr + tpr = 1
    eer = fpr[np.nanargmin(np.absolute((fpr + tpr - 1)))]
    return eer

def compute_eer(loss_file):
    if not os.path.isdir(loss_file):
        loss_file_list = [loss_file]
    else:
        loss_file_list = os.listdir(loss_file)
        loss_file_list = [os.path.join(loss_file, sub_loss_file) for sub_loss_file in loss_file_list]

    optimal_results = RecordResult(auc=np.inf)
    for sub_loss_file in loss_file_list:
        dataset, scores, labels = get_scores_labels(sub_loss_file)
        fpr, tpr, thresholds = metrics.roc_curve(labels, scores, pos_label=0)
        eer = cal_eer(fpr, tpr)

        results = RecordResult(fpr, tpr, eer, dataset, sub_loss_file)

        if optimal_results > results:
            optimal_results = results

        if os.path.isdir(loss_file):
            print(results)
    print('##### optimal result and model = {}'.format(optimal_results))
    return optimal_results

def compute_auc(loss_file):
    if not os.path.isdir(loss_file):
        loss_file_list = [loss_file]
    else:
        loss_file_list = os.listdir(loss_file)
        loss_file_list = [os.path.join(loss_file, sub_loss_file) for sub_loss_file in loss_file_list]

    optimal_results = RecordResult()
    for sub_loss_file in loss_file_list:
        # the name of dataset, loss, and ground truth
        dataset, psnr_records, gt = load_psnr_gt(loss_file=sub_loss_file)

        # the number of videos
        num_videos = len(psnr_records)

        scores = np.array([], dtype=np.float32)
        labels = np.array([], dtype=np.int8)
        # video normalization
        for i in range(num_videos):
            distance = psnr_records[i]

            if NORMALIZE:
                distance -= distance.min()  # distances = (distance - min) / (max - min)
                distance /= distance.max()
                # distance = 1 - distance

            scores = np.concatenate((scores, distance[DECIDABLE_IDX:]), axis=0)
            labels = np.concatenate((labels, gt[i][DECIDABLE_IDX:]), axis=0)
            print('score: ',distance[DECIDABLE_IDX:])
            print('labels: ',gt[i][DECIDABLE_IDX:])
        fpr, tpr, thresholds = metrics.roc_curve(labels, scores, pos_label=0)
        auc = metrics.auc(fpr, tpr)

        result_path = os.getcwd()+'/result_auc/'
        if not os.path.exists(result_path):
            os.makedirs(result_path)

        auroc = {}
        auroc['fpr'] = fpr
        auroc['tpr'] = tpr
        scio.savemat(result_path+'future_pred_auroc.mat',{'auroc':auroc})

        results = RecordResult(fpr, tpr, auc, dataset, sub_loss_file)

        if optimal_results < results:
            optimal_results = results

        if os.path.isdir(loss_file):
            print(results)
    print('##### optimal result and model = {}'.format(optimal_results))
    return optimal_results

def average_psnr(loss_file):
    if not os.path.isdir(loss_file):
        loss_file_list = [loss_file]
    else:
        loss_file_list = os.listdir(loss_file)
        loss_file_list = [os.path.join(loss_file, sub_loss_file) for sub_loss_file in loss_file_list]

    max_avg_psnr = -np.inf
    max_file = ''
    for file in loss_file_list:
        psnr_records = load_psnr(file)

        psnr_records = np.concatenate(psnr_records, axis=0)
        avg_psnr = np.mean(psnr_records)
        if max_avg_psnr < avg_psnr:
            max_avg_psnr = avg_psnr
            max_file = file
        print('{}, average psnr = {}'.format(file, avg_psnr))

    print('max average psnr file = {}, psnr = {}'.format(max_file, max_avg_psnr))

def calculate_psnr(loss_file):
    optical_result = compute_auc(loss_file)
    print('##### optimal result and model = {}'.format(optical_result))

    mean_psnr = []
    for file in os.listdir(loss_file):
        file = os.path.join(loss_file, file)
        dataset, psnr_records, gt = load_psnr_gt(file)

        psnr_records = np.concatenate(psnr_records, axis=0)
        gt = np.concatenate(gt, axis=0)

        mean_normal_psnr = np.mean(psnr_records[gt == 0])
        mean_abnormal_psnr = np.mean(psnr_records[gt == 1])
        mean = np.mean(psnr_records)
        print('mean normal psrn = {}, mean abnormal psrn = {}, mean = {}'.format(
            mean_normal_psnr,
            mean_abnormal_psnr,
            mean)
        )
        mean_psnr.append(mean)
    print('max mean psnr = {}'.format(np.max(mean_psnr)))

def calculate_score(loss_file):
    if not os.path.isdir(loss_file):
        loss_file_path = loss_file
    else:
        optical_result = compute_auc(loss_file)
        loss_file_path = optical_result.loss_file
        print('##### optimal result and model = {}'.format(optical_result))
    dataset, psnr_records, gt = load_psnr_gt(loss_file=loss_file_path)

    # the number of videos
    num_videos = len(psnr_records)

    scores = np.array([], dtype=np.float32)
    labels = np.array([], dtype=np.int8)
    # video normalization
    for i in range(num_videos):
        distance = psnr_records[i]

        distance = (distance - distance.min()) / (distance.max() - distance.min())

        scores = np.concatenate((scores, distance[DECIDABLE_IDX:]), axis=0)
        labels = np.concatenate((labels, gt[i][DECIDABLE_IDX:]), axis=0)

    mean_normal_scores = np.mean(scores[labels == 0])
    mean_abnormal_scores = np.mean(scores[labels == 1])
    print('mean normal scores = {}, mean abnormal scores = {}, '
          'delta = {}'.format(mean_normal_scores, mean_abnormal_scores, mean_normal_scores - mean_abnormal_scores))

def test_func(*args):
    # simulate testing on CUHK AVENUE dataset
    dataset = GroundTruthLoader.AVENUE

    # load the ground truth
    gt_loader = GroundTruthLoader()
    gt = gt_loader(dataset=dataset)

    num_videos = len(gt)

    simulated_results = {
        'dataset': dataset,
        'psnr': []
    }

    simulated_psnr = []
    for i in range(num_videos):
        sub_video_length = gt[i].shape[0]
        simulated_psnr.append(np.random.random(size=sub_video_length))

    simulated_results['psnr'] = simulated_psnr

    # writing to file, 'generated_loss.bin'
    with open('generated_loss.bin', 'wb') as writer:
        pickle.dump(simulated_results, writer, pickle.HIGHEST_PROTOCOL)

    print(file_path.name)
    result = compute_auc(file_path.name)

    print('optimal = {}'.format(result))

eval_type_function = {
    'compute_auc': compute_auc,
    'compute_eer': compute_eer,
    'precision_recall_auc': precision_recall_auc,
    'calculate_psnr': calculate_psnr,
    'calculate_score': calculate_score,
    'average_psnr': average_psnr,
    'average_psnr_sample': average_psnr
}

def evaluate(eval_type, save_file):
    assert eval_type in eval_type_function, 'there is no type of evaluation {}, please check {}' \
        .format(eval_type, eval_type_function.keys())
    eval_func = eval_type_function[eval_type]
    optimal_results = eval_func(save_file)
    return optimal_results

if __name__ == '__main__':
    args = parser_args()

    eval_type = args.type
    file_path = args.file

    print('Evaluate type = {}'.format(eval_type))
    print('File path = {}'.format(file_path))

    if eval_type == 'test_func':
        test_func()
    else:
        evaluate(eval_type, file_path)


import numpy as np
import os
import sys
import esutil
import time

from .fgcmChisq import FgcmChisq

import multiprocessing

from .sharedNumpyMemManager import SharedNumpyMemManager as snmm

class FgcmBrightObs(object):
    """
    Class to compute approximate photometric magnitudes using a
      brightest-observation algorithm

    parameters
    ----------
    fgcmConfig: FgcmConfig
       Config object
    fgcmPars: FgcmParameters
       Parameter object
    fgcmStars: FgcmStars
       Star object
    fgcmLUT: FgcmLUT
       LUT object

    Config variables
    ----------------
    brightObsGrayMax: float
       Maximum gray compared to mean to consider averaging
    nCore: int
       Number of cores to run on (via multiprocessing)
    nStarPerRun: int
       Number of stars per run (too many uses more memory)
    """
    def __init__(self,fgcmConfig,fgcmPars,fgcmStars,fgcmLUT):

        self.fgcmLog = fgcmConfig.fgcmLog

        self.fgcmLog.debug('Initializing FgcmBrightObs')

        # need fgcmPars because it tracks good exposures
        self.fgcmPars = fgcmPars
        # need fgcmStars because it has the stars (duh)
        self.fgcmStars = fgcmStars
        # and fgcmLUT for the SEDs (this makes me unhappy)
        self.fgcmLUT = fgcmLUT

        self.brightObsGrayMax = fgcmConfig.brightObsGrayMax
        self.nCore = fgcmConfig.nCore
        self.nStarPerRun = fgcmConfig.nStarPerRun
        self.bandFitIndex = fgcmConfig.bandFitIndex
        self.quietMode = fgcmConfig.quietMode

        if (fgcmConfig.useSedLUT and self.fgcmLUT.hasSedLUT):
            self.useSedLUT = True
        else:
            self.useSedLUT = False

    def brightestObsMeanMag(self,debug=False,computeSEDSlopes=False):
        """
        Compute the mean magnitude of every object/band from the brightest observations.

        Parameters
        ----------
        debug: bool, default=False
           Debug mode, no multiprocessing
        computeSEDSlopes: bool, default=False
           Compute first-order SED slopes from mean magnitudes
        """

        if (not self.fgcmStars.magStdComputed):
            raise ValueError("Must run FgcmChisq to compute magStd before FgcmBrightObs")

        startTime=time.time()
        if not self.quietMode:
            self.fgcmLog.info('Selecting good stars from Bright Observations')

        self.debug = debug
        self.computeSEDSlopes = computeSEDSlopes

        # reset numbers
        snmm.getArray(self.fgcmStars.objMagStdMeanHandle)[:] = 99.0
        snmm.getArray(self.fgcmStars.objMagStdMeanErrHandle)[:] = 99.0

        goodStars = self.fgcmStars.getGoodStarIndices(includeReserve=True)

        if (goodStars.size == 0):
            raise ValueError("No good stars to fit!")

        # do global pre-matching before giving to workers, because
        #  it is faster this way

        obsObjIDIndex = snmm.getArray(self.fgcmStars.obsObjIDIndexHandle)
        obsExpIndex = snmm.getArray(self.fgcmStars.obsExpIndexHandle)
        obsFlag = snmm.getArray(self.fgcmStars.obsFlagHandle)

        preStartTime=time.time()
        self.fgcmLog.debug('Pre-matching stars and observations...')

        goodStarsSub, goodObs = self.fgcmStars.getGoodObsIndices(goodStars)

        self.fgcmLog.debug('Pre-matching done in %.1f sec.' %
                           (time.time() - preStartTime))

        if (self.debug):
            self._worker((goodStars,goodObs))
        else:
            if not self.quietMode:
                self.fgcmLog.info('Running BrightObs on %d cores' % (self.nCore))

            # split goodStars into a list of arrays of roughly equal size

            prepStartTime = time.time()
            nSections = goodStars.size // self.nStarPerRun + 1
            goodStarsList = np.array_split(goodStars,nSections)

            # is there a better way of getting all the first elements from the list?
            #  note that we need to skip the first which should be zero (checked above)
            #  see also fgcmChisq.py
            # splitValues is the first of the goodStars in each list
            splitValues = np.zeros(nSections-1,dtype='i4')
            for i in range(1,nSections):
                splitValues[i-1] = goodStarsList[i][0]

            # get the indices from the goodStarsSub matched list
            splitIndices = np.searchsorted(goodStars[goodStarsSub], splitValues)

            # and split along these indices
            goodObsList = np.split(goodObs,splitIndices)

            workerList = list(zip(goodStarsList,goodObsList))

            # reverse sort so the longest running go first
            workerList.sort(key=lambda elt:elt[1].size, reverse=True)

            self.fgcmLog.debug('Using %d sections (%.1f seconds)' %
                               (nSections,time.time() - prepStartTime))

            # make a pool
            mp_ctx = multiprocessing.get_context("fork")
            pool = mp_ctx.Pool(processes=self.nCore)
            pool.map(self._worker,workerList,chunksize=1)
            pool.close()
            pool.join()

        if not self.quietMode:
            self.fgcmLog.info('Finished BrightObs in %.2f seconds.' %
                              (time.time() - startTime))

    def _worker(self,goodStarsAndObs):
        """
        Multiprocessing worker for FgcmBrightObs.  Not to be called on its own.

        parameters
        ----------
        goodStarsAndObs: tuple[2]
           (goodStars, goodObs)
        """

        # NOTE: No logging is allowed in the _worker method

        goodStars = goodStarsAndObs[0]
        goodObs = goodStarsAndObs[1]

        objMagStdMean = snmm.getArray(self.fgcmStars.objMagStdMeanHandle)
        objMagStdMeanErr = snmm.getArray(self.fgcmStars.objMagStdMeanErrHandle)
        objNGoodObs = snmm.getArray(self.fgcmStars.objNGoodObsHandle)

        obsExpIndex = snmm.getArray(self.fgcmStars.obsExpIndexHandle)
        obsBandIndex = snmm.getArray(self.fgcmStars.obsBandIndexHandle)
        obsObjIDIndex = snmm.getArray(self.fgcmStars.obsObjIDIndexHandle)
        # obsMagADUErr = snmm.getArray(self.fgcmStars.obsMagADUErrHandle)
        # Note that this will be the same when we don't fit for the error
        obsMagADUModelErr = snmm.getArray(self.fgcmStars.obsMagADUModelErrHandle)
        obsMagStd = snmm.getArray(self.fgcmStars.obsMagStdHandle)
        obsFlag = snmm.getArray(self.fgcmStars.obsFlagHandle)

        # and the arrays for locking access
        objMagStdMeanLock = snmm.getArrayBase(self.fgcmStars.objMagStdMeanHandle).get_lock()

        # and cut to those exposures that are not flagged
        if (self.debug):
            startTime = time.time()
            self.fgcmLog.debug('Cutting to good exposures')

        gd,=np.where((self.fgcmPars.expFlag[obsExpIndex[goodObs]] == 0) &
                     (obsFlag[goodObs] == 0))
        goodObs = goodObs[gd]

        if (self.debug):
            startTime=time.time()
            self.fgcmLog.debug('Cutting to sub-indices.')

        obsObjIDIndexGO = obsObjIDIndex[goodObs]
        obsBandIndexGO = obsBandIndex[goodObs]
        obsMagStdGO = obsMagStd[goodObs]

        if (self.debug):
            self.fgcmLog.debug('Cut to sub-indices in %.1f seconds.' %
                             (time.time() - startTime))

        obsMagErr2GO = obsMagADUModelErr[goodObs]**2.

        # new version using fmin.at()

        # start with the mean temp var, set to 99s.
        objMagStdMeanTemp = np.zeros_like(objMagStdMean)
        objMagStdMeanTemp[:,:] = 99.0

        # find the brightest (minmag) object at each index
        np.fmin.at(objMagStdMeanTemp,
                   (obsObjIDIndexGO, obsBandIndexGO),
                   obsMagStdGO)

        # now which observations are bright *enough* to consider?
        brightEnoughGO, = np.where((obsMagStdGO -
                                    objMagStdMeanTemp[obsObjIDIndexGO,
                                                      obsBandIndexGO]) <=
                                   self.brightObsGrayMax)

        # need to take the weighted mean, so a temp array here
        #  (memory issues?)
        wtSum = np.zeros_like(objMagStdMean,dtype='f8')
        objNGoodObsTemp = np.zeros_like(objNGoodObs)
        objMagStdMeanTemp[:,:] = 0

        obsMagErr2GOBE = obsMagErr2GO[brightEnoughGO]

        np.add.at(wtSum,
                  (obsObjIDIndexGO[brightEnoughGO],
                   obsBandIndexGO[brightEnoughGO]),
                  1./obsMagErr2GOBE)
        np.add.at(objMagStdMeanTemp,
                  (obsObjIDIndexGO[brightEnoughGO],
                   obsBandIndexGO[brightEnoughGO]),
                  obsMagStdGO[brightEnoughGO]/obsMagErr2GOBE)
        np.add.at(objNGoodObsTemp,
                  (obsObjIDIndexGO[brightEnoughGO],
                   obsBandIndexGO[brightEnoughGO]),
                  1)

        # these are good object/bands that were observed
        gd=np.where(wtSum > 0.0)

        # acquire lock to save values
        objMagStdMeanLock.acquire()

        objMagStdMean[gd] = objMagStdMeanTemp[gd] / wtSum[gd]
        objMagStdMeanErr[gd] = np.sqrt(1./wtSum[gd])
        objNGoodObs[gd] = objNGoodObsTemp[gd]

        # and release
        objMagStdMeanLock.release()

        # finally, compute SED slopes if desired
        if (self.computeSEDSlopes):
            if (self.useSedLUT):
                self.fgcmStars.computeObjectSEDSlopesLUT(goodStars, self.fgcmLUT)
            else:
                self.fgcmStars.computeObjectSEDSlopes(goodStars)

        # and we're done

    def __getstate__(self):
        # Don't try to pickle the logger.

        state = self.__dict__.copy()
        del state['fgcmLog']
        return state


import warnings; warnings.simplefilter('ignore')

from ddpg_agent.agent import DDPG#, Q_a_frames_spec
from ddpg_agent.visualizations import plot_quadcopter_episode
from ddpg_agent.quadcopter_environment import Task
import numpy as np
from collections import namedtuple
from hyperopt import STATUS_OK
import json

task = Task(init_pose=np.array([0., 0., 10, 0., 0., 0.]),
            #init_pose=np.array([0., 0., 10, math.pi/2., math.pi/2., 0.]),
            init_velocities=np.array([0., 0., 0.]),
            init_angle_velocities=np.array([0., 0., 0.]),
            runtime=10.,
            vert_dist_thresh=1, horiz_dist_thresh=1,
            target_steps_within_goal=25,
            target_pos=np.array([0., 0., 20.]),
           )

NoiseParams = namedtuple('NoiseParams',['ou_mu','ou_theta','ou_sigma','n_episodes','eps'])

def noise_evaluator(params):
    """ Evaluator to test different noise parameters looking to maximize average training_score.
        No model learning is done during this step.
    """
    params = NoiseParams(**params)
    print(params)
    scores = []
    histories = []
    for i in range(3):
        agent = DDPG(task,
                     ou_mu=params.ou_mu, ou_theta=params.ou_theta, ou_sigma=params.ou_sigma,
                     replay_buffer_size=0,
                     replay_batch_size=params.n_episodes, # suppress model training
                    )
        # agent.print_summary()
        agent.train_n_episodes(params.n_episodes, eps_decay=0, run_tests=False,
                               # Notice we're acting randomly for all n_episodes
                               act_random_first_n_episodes=params.n_episodes,
                               primary_exploration_eps=params.eps, )
        scores.append( np.mean([ep.score for ep in agent.history.training_episodes]) )
        histories.append( dict(agent.history) )
    # Return a negative score (since hyperopt will minimize this function)
    # return -np.mean(scores)
    return {'status':STATUS_OK,
            'loss': -np.mean(scores),
            'loss_variance': np.var(scores),
            'attachments':{
                'history': json.dumps(histories).encode('utf-8'),
                },
            'scores':scores,
            }

LearningParams = namedtuple('LearningParams',[
        'ou_mu','ou_theta','ou_sigma','n_episodes','eps',
        'eps_decay','primary_exploration_eps','act_random_first_n_episodes',
        'discount_factor','replay_buffer_size','replay_batch_size',
        'tau_actor','tau_critic','lr_actor','lr_critic','output_action_regularizer',
    ])

def evaluator(params):
    """ Evaluator to test learning parameters, using the noise parameters learned using
        the evaluator above.
    """
    params = LearningParams(*params)
    print(params)
    scores = []
    histories = []
    for i in range(3):
        agent = DDPG(task, discount_factor=params.discount_factor,
                     ou_mu=params.ou_mu, ou_theta=params.ou_theta, ou_sigma=params.ou_sigma,
                     replay_buffer_size=params.replay_buffer_size,
                     replay_batch_size=params.replay_batch_size,
                     tau_actor=params.tau_actor, tau_critic=params.tau_critic,
                     lr_actor=params.lr_actor, lr_critic=params.lr_critic, #activation_fn_actor='tanh',
        #              relu_alpha_actor=.01, relu_alpha_critic=.01,
        #              l2_reg_actor=.01, l2_reg_critic=.01,
        #              bn_momentum_actor=0, bn_momentum_critic=.7,
        #              q_a_frames_spec=q_a_frames_spec,
        #              do_preprocessing=False,
        #              input_bn_momentum_actor=.7,
        #              input_bn_momentum_critic=.7,
        #              activity_l2_reg=50,
                     output_action_regularizer=params.output_action_regularizer,
                    )
        # agent.print_summary()

        agent.train_n_episodes(params.n_episodes, eps=params.eps, eps_decay=params.eps_decay,
                               act_random_first_n_episodes=params.act_random_first_n_episodes,
                               primary_exploration_eps=params.primary_exploration_eps,
                               )

        # Return a negative score (since hyperopt will minimize this function)
        # Mean test score/variance, last 50 episodes
        last_n_scores = [e.score for e in agent.history.test_episodes[-50:]]
        scores.append(np.mean(last_n_scores))
        histories.append( dict(agent.history) )

        #return -np.mean(last_n_scores)/np.std(last_n_scores)

    return {'status':STATUS_OK,
        'loss': -np.mean(scores),
        'loss_variance': np.var(scores),
        'attachments':{
            'history': json.dumps(histories).encode('utf-8'),
            },
        'scores':scores,
        }

# labels = ['time', 'x', 'y', 'z', 'phi', 'theta', 'psi', 'x_velocity',
#           'y_velocity', 'z_velocity', 'phi_velocity', 'theta_velocity',
#           'psi_velocity', 'rotor_speed1', 'rotor_speed2', 'rotor_speed3', 'rotor_speed4
# [task.sim.time] + list(task.sim.pose) + list(task.sim.v) + list(task.sim.angular_v) + list(rotor_speeds)

# def episode_callback(episode_num):
#     last_training_episode = agent.history.training_episodes[-1]
#     last_test_episode = agent.history.test_episodes[-1]
#     num_steps = last_training_episode.last_step - last_training_episode.first_step+1
#     message = "Episode %i - epsilon: %7.4g, memory: %i, step: %i, training score: %6.2f, test score: %6.2f"\
#                 %(episode_num,
#                   last_training_episode.epsilon,
#                   len(agent.memory),
#                   num_steps,
#                   last_training_episode.score,
#                   last_test_episode.score
#                  )
#     print(message)
#
#     if episode_num%10==0:
#         fig = plot_quadcopter_episode(last_training_episode)
#         display(fig)

# agent.set_episode_callback(episode_callback)

# def max_training_score_callback(episode):
#     last_training_episode = agent.history.training_episodes[-1]
#     print("New best training score.")
#     fig = plot_quadcopter_episode(last_training_episode)
#     display(fig)
#
# # agent.set_max_training_score_callback(max_training_score_callback)
#
# def max_test_score_callback(episode):
#     last_test_episode = agent.history.test_episodes[-1]
#     print("New best test score.")
#     fig = plot_quadcopter_episode(last_test_episode)
#     display(fig)

# agent.set_max_test_score_callback(max_test_score_callback)