import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from utils import utils_sr
import torch
from argparse import ArgumentParser
from utils.utils_restoration import rgb2y, psnr, array2tensor, tensor2array
import sys
from matplotlib.ticker import MaxNLocator

class PnP_restoration():

    def __init__(self, hparams):

        self.hparams = hparams
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.initialize_cuda_denoiser()

    def initialize_cuda_denoiser(self):
        '''
        Initialize the denoiser model with the given pretrained ckpt
        '''
        sys.path.append('../GS_denoising/')
        from lightning_denoiser import GradMatch
        parser2 = ArgumentParser(prog='utils_restoration.py')
        parser2 = GradMatch.add_model_specific_args(parser2)
        parser2 = GradMatch.add_optim_specific_args(parser2)
        hparams = parser2.parse_known_args()[0]
        if 'nb_4' in self.hparams.pretrained_checkpoint :
            hparams.DRUNET_nb = 4
        hparams.grad_matching = self.hparams.grad_matching
        hparams.act_mode = 's'
        self.denoiser_model = GradMatch(hparams)
        checkpoint = torch.load(self.hparams.pretrained_checkpoint, map_location=self.device)
        self.denoiser_model.load_state_dict(checkpoint['state_dict'])
        self.denoiser_model.eval()
        for i, v in self.denoiser_model.named_parameters():
            v.requires_grad = False
        self.denoiser_model = self.denoiser_model.to(self.device)
        if self.hparams.precision == 'double' :
            if self.denoiser_model is not None:
                self.denoiser_model.double()

    def initialize_prox(self, img, degradation):
        '''
        calculus for future prox computatations
        :param img: degraded image
        :param degradation: 2D blur kernel for deblurring and SR, mask for inpainting
        '''
        if self.hparams.degradation_mode == 'deblurring' :
            self.k = degradation
            self.k_tensor = array2tensor(np.expand_dims(self.k, 2)).double().to(self.device)
            self.FB, self.FBC, self.F2B, self.FBFy = utils_sr.pre_calculate_prox(img, self.k_tensor, 1)
        elif self.hparams.degradation_mode == 'SR':
            self.k = degradation
            self.k_tensor = array2tensor(np.expand_dims(self.k, 2)).double().to(self.device)
            self.FB, self.FBC, self.F2B, self.FBFy = utils_sr.pre_calculate_prox(img, self.k_tensor, 2)
        elif self.hparams.degradation_mode == 'inpainting':
            self.M = array2tensor(degradation).double().to(self.device)
            self.My = self.M*img
        else:
            print('degradation mode not treated')

    def calculate_prox(self, img):
        '''
        Calculation of the proximal mapping of the data term f
        :param img: input for the prox
        :return: prox_f(img)
        '''
        if self.hparams.degradation_mode == 'deblurring':
            rho = torch.tensor([1/self.hparams.lamb]).double().repeat(1, 1, 1, 1).to(self.device)
            px = utils_sr.prox_solution(img.double(), self.FB, self.FBC, self.F2B, self.FBFy, rho, 1)
        elif self.hparams.degradation_mode == 'SR':
            rho = torch.tensor([1 /self.hparams.lamb]).double().repeat(1, 1, 1, 1).to(self.device)
            px = utils_sr.prox_solution(img.double(), self.FB, self.FBC, self.F2B, self.FBFy, rho, self.hparams.sf)
        elif self.hparams.degradation_mode == 'inpainting':
            if self.hparams.noise_level_img > 1e-2:
                px = (self.hparams.lamb*self.My + img)/(self.hparams.lamb*self.M+1)
            else :
                px = self.My + (1-self.M)*img
        else:
            print('degradation mode not treated')
        return px

    def calculate_grad(self, img):
        '''
        Calculation of the gradient of the data term f
        :param img: input for the prox
        :return: \nabla_f(img)
        '''
        if self.hparams.degradation_mode == 'deblurring' :
            grad = utils_sr.grad_solution(img.double(), self.FB, self.FBC, self.FBFy, 1)
        if self.hparams.degradation_mode == 'SR' :
            grad = utils_sr.grad_solution(img.double(), self.FB, self.FBC, self.FBFy, self.hparams.sf)
        return grad

    def calculate_regul(self,y,x,g):
        '''
        Calculation of the regularization (1/tau)*phi_sigma(y)
        :param y: Point where to evaluate
        :param x: D^{-1}(y)
        :param g: Precomputed regularization function value at x
        :return: regul(y)
        '''
        regul = (1 / self.hparams.lamb) * (g - (1 / 2) * torch.norm(x - y, p=2) ** 2)
        return regul

    def calulate_data_term(self,y,img):
        '''
        Calculation of the data term value f(y)
        :param y: Point where to evaluate F
        :param img: Degraded image
        :return: f(y)
        '''
        if self.hparams.degradation_mode == 'deblurring':
            deg_y = utils_sr.imfilter(y.double(), self.k_tensor[0].double().flip(1).flip(2).expand(3, -1, -1, -1))
            f = 0.5 * torch.norm(img - deg_y, p=2) ** 2
        elif self.hparams.degradation_mode == 'SR':
            deg_y = utils_sr.imfilter(y.double(), self.k_tensor[0].double().flip(1).flip(2).expand(3, -1, -1, -1))
            deg_y = deg_y[..., 0::self.hparams.sf, 0::self.hparams.sf]
            f = 0.5 * torch.norm(img - deg_y, p=2) ** 2
        elif self.hparams.degradation_mode == 'inpainting':
            deg_y = self.M * y.double()
            f = 0.5 * torch.norm(img - deg_y, p=2) ** 2
        else:
            print('degradation not implemented')
        return f

    def calculate_F(self, y, x, g, img):
        '''
        Calculation of the objective function value f(y) + (1/tau)*phi_sigma(y)
        :param y: Point where to evaluate F
        :param x: D^{-1}(y)
        :param g: Precomputed regularization function value at x
        :param img: Degraded image
        :return: F(y)
        '''
        regul = self.calculate_regul(y,x,g)
        if self.hparams.no_data_term:
            F = regul
        else:
            f = self.calulate_data_term(y,img)
            F = f + regul
        return F.item()

    def calculate_lyapunov_DRS(self,y,z,x,g,img):
        '''
            Calculation of the Lyapunov function value Psi(x)
            :param x: Point where to evaluate F
            :param y,z: DRS iterations initialized at x
            :param g: Precomputed regularization function value at x
            :param img: Degraded image
            :return: Psi(x)
        '''
        regul = self.calculate_regul(y,x,g)
        f = self.calulate_data_term(z, img)
        Psi = regul + f + (1 / self.hparams.lamb) * (torch.sum(torch.mul(y-x,y-z)) + (1/2) * torch.norm(y - z, p=2) ** 2)
        return Psi

    def restore(self, img, init_im, clean_img, degradation,extract_results=False):
        '''
        Compute GS-PnP restoration algorithm
        :param img: Degraded image
        :param clean_img: ground-truth clean image
        :param degradation: 2D blur kernel for deblurring and SR, mask for inpainting
        :param extract_results: Extract information for subsequent image or curve saving
        '''

        if extract_results:
            y_list, z_list, x_list, Dg_list, psnr_tab, g_list, Dx_list, F_list, Psi_list = [], [], [], [], [], [], [], [], []

        i = 0 # iteration counter

        img_tensor = array2tensor(init_im).to(self.device) # for GPU computations (if GPU available)
        self.initialize_prox(img_tensor, degradation)  # prox calculus that can be done outside of the loop

        # Initialization of the algorithm
        if self.hparams.degradation_mode == 'SR':
            x0 = cv2.resize(init_im, (img.shape[1] * self.hparams.sf, img.shape[0] * self.hparams.sf),interpolation=cv2.INTER_CUBIC)
            x0 = utils_sr.shift_pixel(x0, self.hparams.sf)
            x0 = array2tensor(x0).to(self.device)
        else:
            x0 = array2tensor(init_im).to(self.device)

        if extract_results:  # extract np images and PSNR values
            out_x = tensor2array(x0.cpu())
            current_x_psnr = psnr(clean_img, out_x)
            if self.hparams.print_each_step:
                print('current x PSNR : ', current_x_psnr)
            psnr_tab.append(current_x_psnr)
            x_list.append(out_x)

        x = x0

        if self.hparams.use_hard_constraint:
            x = torch.clamp(x, 0, 1)

        # Initialize Lyapunov
        diff_Psi = 1
        Psi_old = 1
        Psi = Psi_old

        while i < self.hparams.maxitr and abs(diff_Psi)/Psi_old > self.hparams.relative_diff_Psi_min:

            if self.hparams.inpainting_init :
                if i < self.hparams.n_init:
                    self.sigma_denoiser = 50
                else :
                    self.sigma_denoiser = self.hparams.sigma_denoiser
            else :
                self.sigma_denoiser = self.hparams.sigma_denoiser

            x_old = x
            Psi_old = Psi

            if self.hparams.PnP_algo == 'PGD':
                # Gradient step
                gradx = self.calculate_grad(x_old)
                z = x_old - self.hparams.lamb*gradx
                # Denoising step
                torch.set_grad_enabled(True)
                Dg, N = self.denoiser_model.calculate_grad(z, self.hparams.sigma_denoiser / 255.)
                torch.set_grad_enabled(False)
                Dg = Dg.detach()
                N = N.detach()
                g = 0.5 * (torch.norm(z.double() - N.double(), p=2) ** 2)
                Dz = z - Dg
                Dx = Dz
                x = (1 - self.hparams.alpha) * z + self.hparams.alpha*Dz
                y = x
                # Hard constraint
                if self.hparams.use_hard_constraint:
                    x = torch.clamp(x,0,1)
                # Calculate Objective
                F = self.calculate_F(x, z, g, img_tensor)

            elif self.hparams.PnP_algo == 'DRS':
                # Denoising step
                torch.set_grad_enabled(True)
                Dg, N = self.denoiser_model.calculate_grad(x_old, self.hparams.sigma_denoiser / 255.)
                torch.set_grad_enabled(False)
                Dg = Dg.detach()
                N = N.detach()
                g = 0.5 * (torch.norm(x_old.double() - N.double(), p=2) ** 2)
                Dx = x_old - Dg
                y = (1 - self.hparams.alpha)*x_old + self.hparams.alpha*Dx
                # Hard constraint
                if self.hparams.use_hard_constraint:
                    y = torch.clamp(y,0,1)
                # Proximal step
                z = self.calculate_prox(2*y-x_old)
                # Calculate Lyapunov
                Psi = self.calculate_lyapunov_DRS(y,z,x,g,img_tensor)
                diff_Psi = Psi-Psi_old
                # Calculate Objective
                F = self.calculate_F(y, x, g, img_tensor)
                # Final step
                x = x_old + (z-y)

            elif self.hparams.PnP_algo == 'DRSdiff':

                # Proximal step
                y = self.calculate_prox(x_old)
                y2 = 2*y-x_old

                # Denoising step
                torch.set_grad_enabled(True)
                Dg, N = self.denoiser_model.calculate_grad(y2, self.hparams.sigma_denoiser / 255.)
                torch.set_grad_enabled(False)
                Dg = Dg.detach()
                N = N.detach()
                g = 0.5 * (torch.norm(y2.double() - N.double(), p=2) ** 2)
                Dx = y2 - Dg
                z = (1 - self.hparams.alpha) * y2 + self.hparams.alpha * Dx
                # Hard constraint
                if self.hparams.use_hard_constraint:
                    z = torch.clamp(z, 0, 1)

                # Calculate Lyapunov
                Psi = self.calculate_lyapunov_DRS(y,z,x,g,img_tensor)
                diff_Psi = Psi-Psi_old
                # Calculate Objective
                F = self.calculate_F(y, x, g, img_tensor)
                # Final step
                x = x_old + (z-y)

            else :
                print('algo not implemented')

            # Logging
            if extract_results:
                out_y = tensor2array(y.cpu())
                out_z = tensor2array(z.cpu())
                out_x = tensor2array(x.cpu())
                current_y_psnr = psnr(clean_img, out_y)
                current_z_psnr = psnr(clean_img, out_z)
                current_x_psnr = psnr(clean_img, out_x)
                if self.hparams.print_each_step:
                    print('iteration : ', i)
                    print('current y PSNR : ', current_y_psnr)
                    print('current z PSNR : ', current_z_psnr)
                    print('current x PSNR : ', current_x_psnr)
                y_list.append(out_y)
                x_list.append(out_x)
                z_list.append(out_z)
                Dx_list.append(tensor2array(Dx.cpu()))
                Dg_list.append(torch.norm(Dg).cpu().item())
                g_list.append(g.cpu().item())
                psnr_tab.append(current_x_psnr)
                F_list.append(F)
                Psi_list.append(Psi)

            # next iteration
            i += 1

        output_img = tensor2array(y.cpu())
        output_psnr = psnr(clean_img, output_img)
        output_psnrY = psnr(rgb2y(clean_img), rgb2y(output_img))

        if extract_results:
            return output_img, output_psnr, output_psnrY, x_list, np.array(z_list), np.array(y_list), np.array(Dg_list), np.array(psnr_tab), np.array(Dx_list), np.array(g_list), np.array(F_list), np.array(Psi_list)
        else:
            return output_img, output_psnr, output_psnrY

    def initialize_curves(self):

        self.conv = []
        self.PSNR = []
        self.g = []
        self.Dg = []
        self.F = []
        self.Psi = []
        self.lip_algo = []
        self.lip_D = []
        self.lip_Dg = []

    def update_curves(self, x_list, Dx_list, psnr_tab, Dg_list, g_list, F_list, Psi_list):

        self.F.append(F_list)
        self.Psi.append(Psi_list)
        self.g.append(g_list)
        self.Dg.append(Dg_list)
        self.PSNR.append(psnr_tab)
        self.conv.append(np.array([(np.linalg.norm(x_list[k + 1] - x_list[k]) ** 2) for k in range(len(x_list) - 1)]) / np.sum(np.abs(x_list[0]) ** 2))
        self.lip_algo.append(np.sqrt(np.array([np.sum(np.abs(x_list[k + 1] - x_list[k]) ** 2) for k in range(1, len(x_list) - 1)]) / np.array([np.sum(np.abs(x_list[k] - x_list[k - 1]) ** 2) for k in range(1, len(x_list[:-1]))])))
        self.lip_D.append(np.sqrt(np.array([np.sum(np.abs(Dx_list[i + 2] - Dx_list[i+1]) ** 2) for i in range(len(Dx_list) - 2)]) / np.array([np.sum(np.abs(x_list[i+1] - x_list[i]) ** 2) for i in range(len(Dx_list) - 2)])))
        self.lip_Dg.append(np.sqrt(np.array([np.sum(np.abs(Dg_list[i + 2] - Dg_list[i+1]) ** 2) for i in range(len(Dx_list) - 2)]) / np.array([np.sum(np.abs(x_list[i+1] - x_list[i]) ** 2) for i in range(len(Dg_list) - 2)])))

    def save_curves(self, save_path):

        import matplotlib
        matplotlib.rcParams.update({'font.size': 10})
        matplotlib.rcParams['lines.linewidth'] = 2
        matplotlib.style.use('seaborn-darkgrid')

        plt.figure(0)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.g)):
            plt.plot(self.g[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'g.png'),bbox_inches="tight")

        plt.figure(1)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.PSNR)):
            plt.plot(self.PSNR[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'PSNR.png'),bbox_inches="tight")

        plt.figure(2)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.F)):
            plt.plot(self.Psi[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'Liapunov.png'), bbox_inches="tight")

        plt.figure(22)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.F)):
            plt.plot(self.F[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'F.png'), bbox_inches="tight")

        plt.figure(4)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.Dg)):
            Ds_norm = [np.linalg.norm(np.array(self.Dg[i][j])) for j in range(len(self.Dg[i]))]
            plt.plot(Ds_norm, linewidth=1.5, markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'Dg.png'), bbox_inches="tight")

        plt.figure(5)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.conv)):
            plt.plot(self.conv[i], '-o', markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'conv_log.png'), bbox_inches="tight")

        self.conv2 = [[np.min(self.conv[i][:k]) for k in range(1, len(self.conv[i]))] for i in range(len(self.conv))]
        conv_rate = [self.conv2[i][0]*np.array([(1/k) for k in range(1,len(self.conv2[i]))]) for i in range(len(self.conv2))]
        plt.figure(6)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.conv)):
            plt.plot(self.conv2[i], '-', markevery=10)
        plt.plot(conv_rate[i], '--', color='red', label=r'$\mathcal{O}(\frac{1}{K})$')
        plt.semilogy()
        plt.legend()
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'conv_log2.png'), bbox_inches="tight")

        self.conv_sum = [[np.sum(self.conv[i][:k]) for k in range(1, len(self.conv[i]))] for i in range(len(self.conv))]

        plt.figure(7)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.conv_sum)):
            plt.plot(self.conv_sum[i], markevery=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.legend()
        plt.savefig(os.path.join(save_path, 'conv_log_sum.png'), bbox_inches="tight")

        plt.figure(8)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.lip_algo)):
            plt.plot(self.lip_algo[i], '-o', markersize=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'lip_algo.png'))

        plt.figure(9)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.lip_D)):
            plt.plot(self.lip_D[i], '-o', markersize=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'lip_D.png'))

        plt.figure(10)
        fig, ax = plt.subplots()
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        for i in range(len(self.lip_Dg)):
            plt.plot(self.lip_Dg[i], '-o', markersize=10)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        plt.savefig(os.path.join(save_path, 'lip_Dg.png'))

    def add_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument('--dataset_path', type=str, default='../datasets')
        parser.add_argument('--pretrained_checkpoint', type=str,default='../GS_denoising/ckpts/Prox_DRUNet.ckpt')
        parser.add_argument('--PnP_algo', type=str, default='PGD')
        parser.add_argument('--dataset_name', type=str, default='CBSD68')
        parser.add_argument('--sigma_denoiser', type=float)
        parser.add_argument('--sigma_k_denoiser', type=float)
        parser.add_argument('--noise_level_img', type=float, default=2.55)
        parser.add_argument('--maxitr', type=int, default=1000)
        parser.add_argument('--alpha', type=float, default=1)
        parser.add_argument('--lamb', type=float)
        parser.add_argument('--n_images', type=int, default=68)
        parser.add_argument('--relative_diff_Psi_min', type=float, default=1e-8)
        parser.add_argument('--inpainting_init', dest='inpainting_init', action='store_true')
        parser.set_defaults(inpainting_init=False)
        parser.add_argument('--precision', type=str, default='simple')
        parser.add_argument('--n_init', type=int, default=10)
        parser.add_argument('--patch_size', type=int, default=256)
        parser.add_argument('--extract_curves', dest='extract_curves', action='store_true')
        parser.set_defaults(extract_curves=False)
        parser.add_argument('--extract_images', dest='extract_images', action='store_true')
        parser.set_defaults(extract_images=False)
        parser.add_argument('--print_each_step', dest='print_each_step', action='store_true')
        parser.set_defaults(print_each_step=False)
        parser.add_argument('--no_grad_matching', dest='grad_matching', action='store_false')
        parser.set_defaults(grad_matching=True)
        parser.add_argument('--no_data_term', dest='no_data_term', action='store_true')
        parser.set_defaults(no_data_term=False)
        parser.add_argument('--use_hard_constraint', dest='use_hard_constraint', action='store_true')
        parser.set_defaults(no_data_term=False)
        return parser


import sys

import cv2
from keras.models import load_model
import numpy as np

from utils.datasets import get_labels
from utils.inference import detect_faces
from utils.inference import draw_text
from utils.inference import draw_bounding_box
from utils.inference import apply_offsets
from utils.inference import load_detection_model
from utils.inference import load_image
from utils.preprocessor import preprocess_input

# parameters for loading data and images
image_path = sys.argv[1]
detection_model_path = '../trained_models/detection_models/haarcascade_frontalface_default.xml'
emotion_model_path = '../trained_models/emotion_models/fer2013_mini_XCEPTION.54-0.65.hdf5'
emotion_labels = get_labels('fer2013')
font = cv2.FONT_HERSHEY_SIMPLEX

# hyper-parameters for bounding boxes shape
emotion_offsets = (20, 40)
emotion_offsets = (0, 0)

# loading models
face_detection = load_detection_model(detection_model_path)
emotion_classifier = load_model(emotion_model_path, compile=False)

# getting input model shapes for inference
emotion_target_size = emotion_classifier.input_shape[1:3]

# loading images
rgb_image = load_image(image_path, grayscale=False)
gray_image = load_image(image_path, grayscale=True)
gray_image = np.squeeze(gray_image)
gray_image = gray_image.astype('uint8')

faces = detect_faces(face_detection, gray_image)
for face_coordinates in faces:
    x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
    gray_face = gray_image[y1:y2, x1:x2]

    try:
        gray_face = cv2.resize(gray_face, (emotion_target_size))
    except:
        continue

    gray_face = preprocess_input(gray_face, True)
    gray_face = np.expand_dims(gray_face, 0)
    gray_face = np.expand_dims(gray_face, -1)
    emotion_label_arg = np.argmax(emotion_classifier.predict(gray_face))
    emotion_text = emotion_labels[emotion_label_arg]
    color = (0, 0, 255)
    draw_bounding_box(face_coordinates, rgb_image, color)
    draw_text(face_coordinates, rgb_image, emotion_text, color, 0, -50, 1, 2)

bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
import time

# cv2.imwrite('../images/result/'+str(round(time.time()*100))+'predicted_test_image.png', bgr_image)

cv2.imwrite('../images/predicted_test_image.png', bgr_image)


# -*- coding: utf-8 -*-
import sys
sys.path.append("..")
sys.path.append(".")
sys.path.append("./play")

from math import sqrt
from sys import stderr
import time
from board import *
import model
import numpy as np
# noinspection PyUnresolvedReferences
import tensorflow.compat.v1 as tf

tf.disable_v2_behavior()
# gpus = tf.config.list_physical_devices('GPU')
# tf.config.experimental.set_memory_growth(gpus[0], True)

# 限制节点
max_node_cnt = 2 ** 14  # 16384
expand_cnt = 8

class Node(object):
    """
    MCTS的树节点
    """

    def __init__(self):
        self.init_branch()
        self.branch_cnt = None
        self.total_value = None
        self.total_cnt = None
        self.hash = None
        self.move_cnt = None
        self.clear()

    def clear(self):
        self.branch_cnt = 0
        self.total_value = 0.0
        self.total_cnt = 0
        self.hash = 0
        self.move_cnt = -1

    def init_branch(self):
        # 动作的状态空间
        self.move = np.full(BVCNT + 1, VNULL)
        # 先验概率
        self.prob = np.full(BVCNT + 1, 0.0)
        # 价值
        self.value = np.full(BVCNT + 1, 0.0)
        # 赢的价值
        self.value_win = np.full(BVCNT + 1, 0.0)
        # 被访过的次数
        self.visit_cnt = np.full(BVCNT + 1, 0)
        # 配合Board的id属性一起使用==>子结点
        self.next_id = np.full(BVCNT + 1, -1)
        # dtype is required only on Windows.
        self.next_hash = np.full(BVCNT + 1, -1, dtype=np.int64)
        # 评估价值
        self.evaluated = np.full(BVCNT + 1, False)

class Tree(object):
    cp = 2.0
    # 是否到时间停止
    stop = False

    def __init__(self, ckpt_path="model.ckpt", use_gpu=True):
        self.set_sess(ckpt_path, use_gpu)
        self.node = [Node() for _ in range(max_node_cnt)]
        self.main_time = 0.0
        self.byoyomi = 1.0  # 读秒
        self.clear()

    def clear(self):
        self.left_time = self.main_time
        for nd in self.node:
            nd.clear()
        self.node_cnt = 0
        self.root_id = 0
        self.root_move_cnt = 0
        self.node_hashs = {}
        self.eval_cnt = 0
        Tree.stop = False

    def set_sess(self, ckpt_path, use_gpu=True):
        """
        声明模型，开启tf会话
        """
        device_name = "gpu" if use_gpu else "cpu"
        # 使用默认图、显卡设置
        with tf.get_default_graph().as_default(), tf.device("/%s:0" % device_name):
            dn = model.DualNetwork()
            self.x = tf.placeholder("float", shape=[None, BVCNT, FEATURE_CNT])
            self.pv = dn.model(self.x, temp=0.7, dr=1.0)  # return policy, value
            self.sess = dn.create_sess(ckpt_path)

    def evaluate(self, b: Board):
        """
        判断当前局势
        :param b: board()
        :return: policy, value
        """
        return self.sess.run(self.pv,
                             feed_dict={self.x: np.reshape(b.feature, (1, BVCNT, 7))})

    # 删除节点，如果还没有到一半的节点则不删
    def delete_node(self):
        if self.node_cnt < max_node_cnt * 0.5:
            return
        for i in range(max_node_cnt):
            mc = self.node[i].move_cnt
            # 删除已走步数之间的节点
            if 0 <= mc < self.root_move_cnt:
                if self.node[i].hash in self.node_hashs:
                    self.node_hashs.pop(self.node[i].hash)
                self.node[i].clear()

    # 找到当前搜索的根节点
    def create_node(self, b_info, prob):
        hs = b_info[0]

        # 搜索过相同状态则返回该状态
        if hs in self.node_hashs and \
                self.node[self.node_hashs[hs]].hash == hs and \
                self.node[self.node_hashs[hs]].move_cnt == b_info[1]:
            return self.node_hashs[hs]

        # 设置结点标号
        node_id = hs % max_node_cnt
        # 将节点标号映射回零域
        while self.node[node_id].move_cnt != -1:
            node_id = node_id + 1 if node_id + 1 < max_node_cnt else 0

        self.node_hashs[hs] = node_id
        self.node_cnt += 1

        nd = self.node[node_id]
        nd.clear()
        nd.move_cnt = b_info[1]
        nd.hash = hs
        nd.init_branch()
        # 把所有的策略标号加入
        order_ = np.argsort(prob)[::-1]
        for rv in order_:
            # 如果可行，加入子树
            if rv in b_info[2]:
                nd.move[nd.branch_cnt] = rv2ev(rv)
                nd.prob[nd.branch_cnt] = prob[rv]
                nd.branch_cnt += 1

        return node_id

    def  search_branch(self, b:Board, node_id, route):
        nd = self.node[node_id]

        nd_rate = 0.0 if nd.total_cnt == 0 else nd.total_value / nd.total_cnt
        cpsv = Tree.cp * sqrt(nd.total_cnt)

        with np.errstate(divide='ignore', invalid='ignore'):
            rate = nd.value_win / nd.visit_cnt  # including dividing by 0
            rate[~np.isfinite(rate)] = nd_rate  # convert nan, inf to nd_rate
        # 找到最应该访问的下一个节点
        action_value = rate + cpsv * nd.prob / (nd.visit_cnt + 1)
        best = np.argmax(action_value[:nd.branch_cnt])
        # 加入路径
        route.append((node_id, best))
        next_id = nd.next_id[best]
        next_move = nd.move[best]
        # 是否应该继续行动
        head_node = not self.has_next(node_id, best, b.move_cnt + 1) or \
                    nd.visit_cnt[best] < expand_cnt or \
                    (b.move_cnt > BVCNT * 2) or \
                    (next_move == PASS and b.prev_move == PASS)

        b.play(next_move, False)

        if head_node:
            if nd.evaluated[best]:
                value = nd.value[best]
            else:
                prob_, value_ = self.evaluate(b)
                self.eval_cnt += 1
                value = -value_[0]
                nd.value[best] = value
                nd.evaluated[best] = True

                if self.node_cnt > 0.85 * max_node_cnt:
                    self.delete_node()

                next_id = self.create_node(b.info(), prob_[0])
                next_nd = self.node[next_id]
                nd.next_id[best] = next_id
                nd.next_hash[best] = b.hash()

                next_nd.total_value -= nd.value_win[best]
                next_nd.total_cnt += nd.visit_cnt[best]
        else:
            value = -self.search_branch(b, next_id, route)

        nd.total_value += value
        nd.total_cnt += 1
        nd.value_win[best] += value
        nd.visit_cnt[best] += 1

        return value

    # 搜索
    def search(self, b: Board, time_, ponder=False, clean=False):
        start = time.time()
        # 通过模型得到策略
        prob, v = self.evaluate(b)  # 返回policy和value
        self.root_id = self.create_node(b.info(), prob[0])
        self.root_move_cnt = b.move_cnt
        Tree.cp = 0.01 if b.move_cnt < 8 else 1.5

        nd = self.node[self.root_id]
        # 无后分支-》返回pass
        if nd.branch_cnt <= 1:
            stderr.write("\nmove count=%d:\n" % (b.move_cnt + 1))
            self.print_info(self.root_id)
            return PASS, 0.5

        self.delete_node()
        # 找到最好和次好策略
        order_ = np.argsort(nd.visit_cnt[:nd.branch_cnt])[::-1]
        best, second = tuple(order_[:2].tolist())

        win_rate = self.branch_rate(nd, best)

        #         if not ponder and self.byoyomi == 0 and self.left_time < 10:
        #             if nd.visit_cnt[best] < 1000:
        #                 return rv2ev(np.argmax(prob)), 0.5
        #             else:
        #                 stderr.write("\nmove count=%d:\n" % (b.move_cnt + 1))
        #                 self.print_info(self.root_id)
        #                 return nd.move[best], win_rate
        # ￥可调节
        # 进入分支次数并且进最好分支次数明显大于次好
        stand_out = nd.total_cnt > 5000 and nd.visit_cnt[best] > nd.visit_cnt[second] * 100
        # 节点胜率极端
        almost_win = nd.total_cnt > 5000 and (win_rate < 0.1 or win_rate > 0.9)
        # 两个条件均不满足进入思考模式
        if ponder or not (stand_out or almost_win):
            if time_ == 0:
                if self.main_time == 0 or self.left_time < self.byoyomi * 2:
                    time_ = max(self.byoyomi, 1.0)
                else:
                    time_ = self.left_time / (55.0 + max(50 - b.move_cnt, 0.0))

            # search
            search_idx = 1
            self.eval_cnt = 0
            b_cpy = Board()
            while 1:
                # 复制棋盘
                b.copy(b_cpy)
                route = []

                self.search_branch(b_cpy, self.root_id, route)
                search_idx += 1
                if search_idx % 64 == 0:
                    # 满足“思考”和超时其一条件，退出搜索
                    if (ponder and Tree.stop) or time.time() - start > time_:
                        Tree.stop = False
                        break
            # 选出最优着
            order_ = np.argsort(nd.visit_cnt[:nd.branch_cnt])[::-1]
            best, second = tuple(order_[:2].tolist())

        next_move = nd.move[best]
        win_rate = self.branch_rate(nd, best)

        # 如果下一着为pass且有“价值机会，则取‘争取价值’”
        if clean and next_move == PASS:
            if nd.value_win[best] * nd.value_win[second] > 0:
                next_move = nd.move[second]
                win_rate = self.branch_rate(nd, second)
        # 无思考则输出日志
        if not ponder:
            stderr.write("\nmove count=%d: left time=%.1f[sec] evaluated=%d\n" % (
                b.move_cnt + 1, max(self.left_time - time_, 0), self.eval_cnt))
            self.print_info(self.root_id)
            self.left_time = max(0.0, self.left_time - (time.time() - start))

        return next_move, win_rate

    # 判断是否存在下一个节点
    def has_next(self, node_id, br_id, move_cnt):
        nd = self.node[node_id]
        next_id = nd.next_id[br_id]
        return next_id >= 0 and \
               nd.next_hash[br_id] == self.node[next_id].hash and \
               self.node[next_id].move_cnt == move_cnt

    # 进入该分支的概率，胜率正相关，访问分支次数负相关
    def branch_rate(self, nd, id):
        return nd.value_win[id] / max(nd.visit_cnt[id], 1) / 2 + 0.5

    def best_sequence(self, node_id, head_move):
        seq_str = "%-3s" % ev2str(head_move)
        next_move = head_move

        for i in range(7):
            nd = self.node[node_id]
            if next_move == PASS or nd.branch_cnt < 1:
                break

            best = np.argmax(nd.visit_cnt[:nd.branch_cnt])
            if nd.visit_cnt[best] == 0:
                break
            next_move = nd.move[best]
            seq_str += "->%-3s" % ev2str(next_move)

            if not self.has_next(node_id, best, nd.move_cnt + 1):
                break
            node_id = nd.next_id[best]

        return seq_str

    def print_info(self, node_id):
        nd = self.node[node_id]
        order_ = np.argsort(nd.visit_cnt[:nd.branch_cnt])[::-1]
        stderr.write("|move|count  |rate |value|prob | best sequence\n")

        for i in range(min(len(order_), 9)):
            m = order_[i]
            visit_cnt = nd.visit_cnt[m]
            if visit_cnt == 0:
                break

            rate = 0.0 if visit_cnt == 0 else self.branch_rate(nd, m) * 100
            value = (nd.value[m] / 2 + 0.5) * 100

            stderr.write("|%-4s|%7d|%5.1f|%5.1f|%5.1f| %s\n" % (
                ev2str(nd.move[m]), visit_cnt, rate, value, nd.prob[m] * 100,
                self.best_sequence(nd.next_id[m], nd.move[m])))

if __name__ == "__main__":
    b = Board()
    tree = Tree("model.ckpt", use_gpu=False)
    move, _rate = tree.search(b, 3, ponder=True)


import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import normal_init, kaiming_init, constant_init
import math
import numpy as np

from mmdet.ops import ModulatedDeformConvPack, soft_nms
from mmdet.core import multi_apply, force_fp32
from mmdet.models.losses import ct_focal_loss, giou_loss
from mmdet.models.utils import (
    build_conv_layer, build_norm_layer, bias_init_with_prob, ConvModule)
from ..registry import HEADS
from .anchor_head import AnchorHead

class UpsamplingLayers(nn.Sequential):

    def __init__(self, in_channels, out_channels, norm_cfg=dict(type='BN'), no_upsample=False):
        mdcn = ModulatedDeformConvPack(
            in_channels,
            out_channels,
            3,
            stride=1,
            padding=1,
            dilation=1,
            deformable_groups=1)
        layers = []
        layers.append(mdcn)
        if norm_cfg:
            layers.append(build_norm_layer(norm_cfg, out_channels)[1])
        layers.append(nn.ReLU(inplace=True))
        if not no_upsample:
            layers.append(nn.UpsamplingBilinear2d(scale_factor=2))
        super(UpsamplingLayers, self).__init__(*layers)

class ShortcutConnection(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_sizes, conv_cfg):
        super(ShortcutConnection, self).__init__()
        layers = []
        for i, kernel_size in enumerate(kernel_sizes):
            inc = in_channels if i == 0 else out_channels
            padding = (kernel_size - 1) // 2
            if conv_cfg:
                layers.append(
                    build_conv_layer(conv_cfg, inc, out_channels, kernel_size,
                                     padding=padding))
            else:
                layers.append(
                    nn.Conv2d(inc, out_channels, kernel_size, padding=padding))
            if i < len(kernel_sizes) - 1:
                layers.append(nn.ReLU(inplace=True))

        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)

@HEADS.register_module
class TTFLevelHead(AnchorHead):

    def __init__(self,
                 inplanes=(64, 128, 256, 512),
                 planes=(256, 128, 64),
                 down_ratio_b1=8,
                 down_ratio_b2=4,
                 hm_head_channels=256,
                 wh_head_channels=(64, 64),
                 hm_head_conv_num=2,
                 wh_head_conv_num=(2, 2),
                 num_classes=81,
                 shortcut_cfg=(1, 2, 3),
                 wh_scale_factor_b1=16.,
                 wh_scale_factor_b2=16.,
                 alpha=0.54,
                 beta=0.54,
                 hm_weight=1.,
                 wh_weight_b1=5.,
                 wh_weight_b2=5.,
                 b1_min_length=32,
                 b2_max_length=64,
                 level_base_area=True,
                 inf_branch=['b1', 'b2'],
                 use_simple_nms=True,
                 focal_loss_beta=4,
                 focal_b2_only=False,
                 shortcut_conv_cfg=None,
                 head_conv_cfg=None,
                 inf_branch_filter=False,
                 max_objs=128,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN')):
        super(AnchorHead, self).__init__()
        assert len(inplanes) == 4 and len(planes) == 3 and len(shortcut_cfg) == 3
        self.inplanes = inplanes
        self.planes = planes
        self.down_ratio_b1 = down_ratio_b1
        self.down_ratio_b2 = down_ratio_b2
        self.hm_head_channels = hm_head_channels
        self.wh_head_channels = wh_head_channels
        self.hm_head_conv_num = hm_head_conv_num
        self.wh_head_conv_num = wh_head_conv_num
        self.num_classes = num_classes
        self.num_fg = num_classes - 1
        self.shortcut_cfg = shortcut_cfg
        self.wh_scale_factor_b1 = wh_scale_factor_b1
        self.wh_scale_factor_b2 = wh_scale_factor_b2
        self.alpha = alpha
        self.beta = beta
        self.max_objs = max_objs
        self.wh_weight_b1 = wh_weight_b1
        self.wh_weight_b2 = wh_weight_b2
        self.b1_min_length = b1_min_length
        self.b2_max_length = b2_max_length
        self.level_base_area = level_base_area
        self.inf_branch = inf_branch
        self.use_simple_nms = use_simple_nms
        self.focal_loss_beta = focal_loss_beta
        self.focal_b2_only = focal_b2_only
        self.shortcut_conv_cfg = shortcut_conv_cfg
        self.head_conv_cfg = head_conv_cfg
        self.inf_branch_filter = inf_branch_filter
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.base_loc_b1 = None
        self.base_loc_b2 = None
        self.fp16_enabled = False

        self._init_layers()

    def _init_branch_layers(self, planes):
        wh_layers, wh2_layers, hm_layers = [], [], []
        inp = planes
        for i in range(self.wh_head_conv_num[0]):
            wh_layers.append(
                ConvModule(
                    inp,
                    self.wh_head_channels[0],
                    3,
                    padding=1,
                    conv_cfg=self.conv_cfg))
            inp = self.wh_head_channels[0]
        if self.head_conv_cfg:
            wh_layers.append(
                build_conv_layer(
                    self.head_conv_cfg,
                    self.wh_head_channels[0],
                    4,
                    kernel_size=3,
                    padding=1
                )
            )
        else:
            wh_layers.append(nn.Conv2d(self.wh_head_channels[0], 4, 3, padding=1))

        inp = planes
        for i in range(self.wh_head_conv_num[1]):
            wh2_layers.append(
                ConvModule(
                    inp,
                    self.wh_head_channels[1],
                    3,
                    padding=1,
                    conv_cfg=self.conv_cfg))
            inp = self.wh_head_channels[1]
        if self.head_conv_cfg:
            wh2_layers.append(
                build_conv_layer(
                    self.head_conv_cfg,
                    self.wh_head_channels[1],
                    4,
                    kernel_size=3,
                    padding=1
                )
            )
        else:
            wh2_layers.append(nn.Conv2d(self.wh_head_channels[1], 4, 3, padding=1))

        inp = planes
        for i in range(self.hm_head_conv_num):
            hm_layers.append(
                ConvModule(
                    inp,
                    self.hm_head_channels,
                    3,
                    padding=1,
                    conv_cfg=self.conv_cfg))
            inp = self.hm_head_channels
        if self.head_conv_cfg:
            hm_layers.append(
                build_conv_layer(
                    self.head_conv_cfg,
                    self.hm_head_channels,
                    self.num_fg,
                    kernel_size=3,
                    padding=1
                )
            )
        else:
            hm_layers.append(nn.Conv2d(self.hm_head_channels, self.num_fg, 3, padding=1))

        wh_layers = nn.Sequential(*wh_layers)
        wh2_layers = nn.Sequential(*wh2_layers)
        hm_layers = nn.Sequential(*hm_layers)
        return wh_layers, wh2_layers, hm_layers

    def _init_layers(self):
        self.upsample_layers = nn.ModuleList([
            UpsamplingLayers(
                self.inplanes[-1], self.planes[0], norm_cfg=self.norm_cfg),
            UpsamplingLayers(
                self.planes[0], self.planes[1], norm_cfg=self.norm_cfg),
            UpsamplingLayers(
                self.planes[1], self.planes[2], norm_cfg=self.norm_cfg)
        ])

        self.shortcut_layers = nn.ModuleList()
        for (inp, outp, layer_num) in zip(self.inplanes[::-1][1:],
                                          self.planes, self.shortcut_cfg):
            assert layer_num > 0, "Shortcut connection must be included."
            self.shortcut_layers.append(
                ShortcutConnection(inp, outp, [3] * layer_num, self.shortcut_conv_cfg))

        self.wh_b1, self.wh_b2, self.hm = self._init_branch_layers(self.planes[-1])

    def init_weights(self):
        for m in self.upsample_layers.modules():
            if isinstance(m, nn.BatchNorm2d):
                constant_init(m, 1)

        for m in self.shortcut_layers.modules():
            if isinstance(m, nn.Conv2d):
                kaiming_init(m)

        bias_cls = bias_init_with_prob(0.01)
        for m in self.hm.modules():
            if isinstance(m, nn.Conv2d):
                normal_init(m, std=0.01)
        normal_init(self.hm[-1], std=0.01, bias=bias_cls)

        for wh in [self.wh_b1, self.wh_b2]:
            for m in wh.modules():
                if isinstance(m, nn.Conv2d):
                    normal_init(m, std=0.001)

        for m in self.modules():
            if isinstance(m, ModulatedDeformConvPack):
                constant_init(m.conv_offset, 0)

    def forward(self, feats):
        """

        Args:
            feats: list(tensor).

        Returns:
            hm: list(tensor), (batch, 80, h, w).
            wh: list(tensor), (batch, 4, h, w).
        """
        y, shortcuts = [], []
        x = feats[-1]
        for i, shortcut_layer in enumerate(self.shortcut_layers):
            shortcuts.append(shortcut_layer(feats[-i - 2]))

        for i, upsampling_layer in enumerate(self.upsample_layers):
            x = upsampling_layer(x)
            x = x + shortcuts[i]
            y.append(x)

        y_s4 = y[-1]
        hm = self.hm(y_s4)
        wh_b1 = F.relu(self.wh_b1(y_s4)) * self.wh_scale_factor_b1
        wh_b2 = F.relu(self.wh_b2(y_s4)) * self.wh_scale_factor_b2

        return hm, wh_b1, wh_b2

    def get_bboxes_single(self,
                          pred_hm,
                          pred_wh,
                          down_ratio,
                          topk,
                          idx=0):
        batch, cat, height, width = pred_hm.size()
        pred_hm = pred_hm.detach().sigmoid_()
        wh = pred_wh.detach()

        # used maxpool to filter the max score
        heat = self.simple_nms(pred_hm)

        # (batch, topk)
        scores, inds, clses, ys, xs = self._topk(heat, topk=topk)
        xs = xs.view(batch, topk, 1) * down_ratio
        ys = ys.view(batch, topk, 1) * down_ratio

        wh = wh.permute(0, 2, 3, 1).contiguous()
        wh = wh.view(wh.size(0), -1, wh.size(3))
        inds = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), wh.size(2))
        wh = wh.gather(1, inds)

        wh_filter = wh.new_ones((batch, topk), dtype=torch.bool)
        if self.inf_branch_filter:
            area = (wh[..., 2] + wh[..., 0] + 1) * (wh[..., 3] + wh[..., 1] + 1)
            if idx == 0:
                wh_filter = area >= self.b1_min_length ** 2 / 2
            elif idx == 1:
                wh_filter = area <= self.b2_max_length ** 2 * 2

        wh = wh.view(batch, topk, 4)
        clses = clses.view(batch, topk, 1).float()
        scores = scores.view(batch, topk, 1)

        bboxes = torch.cat([
            xs - wh[..., [0]], ys - wh[..., [1]], xs + wh[..., [2]],
            ys + wh[..., [3]]
        ], dim=2)
        return heat, inds, clses, scores, bboxes, xs, ys, wh_filter

    @force_fp32(apply_to=('pred_hm_b1', 'pred_hm_b2', 'pred_wh_b1', 'pred_wh_b2'))
    def get_bboxes(self,
                   pred_hm,
                   pred_wh_b1,
                   pred_wh_b2,
                   img_metas,
                   cfg,
                   rescale=False):
        topk = getattr(cfg, 'max_per_img', 100)
        heat_b1, inds_b1, clses_b1, scores_b1, bboxes_b1, xs_b1, ys_b1, wh_filter_b1 = \
            self.get_bboxes_single(pred_hm_b1, pred_wh_b1, self.down_ratio_b1, topk, idx=0)
        heat_b2, inds_b2, clses_b2, scores_b2, bboxes_b2, xs_b2, ys_b2, wh_filter_b2 = \
            self.get_bboxes_single(pred_hm_b2, pred_wh_b2, self.down_ratio_b2, topk, idx=1)

        # heat_b2 = nn.functional.max_pool2d(heat_b2, (3, 3), stride=1, padding=1)
        # heat_b2 = nn.functional.max_pool2d(heat_b2, (2, 2), stride=2)
        # collide_heat = ((heat_b1 > 0) & (heat_b2 > 0)).max(1)[0]
        # collide_heat = collide_heat.view(collide_heat.size(0), -1)  # (batch, h * w)
        # collide_b1 = collide_heat.gather(-1, )

        result_list = []
        score_thr = getattr(cfg, 'score_thr', 0.01)
        if 'b2' not in self.inf_branch:
            bboxes = bboxes_b1
            scores = scores_b1
            clses = clses_b1
            wh_filter = wh_filter_b1
        elif 'b1' not in self.inf_branch:
            bboxes = bboxes_b2
            scores = scores_b2
            clses = clses_b2
            wh_filter = wh_filter_b2
        else:
            bboxes = torch.cat([bboxes_b1, bboxes_b2], dim=1)
            scores = torch.cat([scores_b1, scores_b2], dim=1)
            clses = torch.cat([clses_b1, clses_b2], dim=1)
            wh_filter = torch.cat([wh_filter_b1, wh_filter_b2], dim=1)

        for batch_i in range(bboxes.shape[0]):
            scores_per_img = scores[batch_i]
            wh_filter_per_img = wh_filter[batch_i]
            scores_keep = (scores_per_img > score_thr).squeeze(-1) & wh_filter_per_img

            scores_per_img = scores_per_img[scores_keep]
            bboxes_per_img = bboxes[batch_i][scores_keep]
            labels_per_img = clses[batch_i][scores_keep].squeeze(-1)
            img_shape = img_metas[batch_i]['pad_shape']
            bboxes_per_img[:, 0::2] = bboxes_per_img[:, 0::2].clamp(
                min=0, max=img_shape[1] - 1)
            bboxes_per_img[:, 1::2] = bboxes_per_img[:, 1::2].clamp(
                min=0, max=img_shape[0] - 1)

            if rescale:
                scale_factor = img_metas[batch_i]['scale_factor']
                bboxes_per_img /= bboxes_per_img.new_tensor(scale_factor)

            if self.use_simple_nms:
                bboxes_per_img = torch.cat([bboxes_per_img, scores_per_img], dim=1)
            else:
                labels_int_flatten = labels_per_img.int()
                unique_cls_ids = list(set(list(labels_int_flatten.cpu().numpy())))
                bboxes_per_img_per_cls = bboxes_per_img.new_zeros((0, 5))
                labels_per_img_per_cls = labels_int_flatten.new_zeros((0,))
                for cls_id in unique_cls_ids:
                    cls_id_idx = (labels_int_flatten == cls_id)
                    soft_bboxes, ori_idx = soft_nms(torch.cat((
                        bboxes_per_img[cls_id_idx], scores_per_img[cls_id_idx]), dim=1),
                        iou_thr=0.6)
                    unique_labels = labels_int_flatten[cls_id_idx][ori_idx]
                    bboxes_per_img_per_cls = torch.cat((bboxes_per_img_per_cls, soft_bboxes), dim=0)
                    labels_per_img_per_cls = torch.cat((labels_per_img_per_cls, unique_labels))
                bboxes_per_img = bboxes_per_img_per_cls
                labels_per_img = labels_per_img_per_cls.float()

            labels_per_img = labels_per_img.squeeze(-1)
            result_list.append((bboxes_per_img, labels_per_img))

        return result_list

    def loss_single(self,
                    pred_hm,
                    pred_wh,
                    heatmap,
                    box_target,
                    wh_weight,
                    down_ratio,
                    base_loc_name,
                    hm_weight_factor,
                    wh_weight_factor,
                    focal_loss_beta):
        H, W = pred_hm.shape[2:]
        pred_hm = torch.clamp(pred_hm.sigmoid_(), min=1e-4, max=1 - 1e-4)
        loss_cls = ct_focal_loss(pred_hm, heatmap,
                                 beta=focal_loss_beta) * hm_weight_factor

        if getattr(self, base_loc_name) is None or H != getattr(self, base_loc_name).shape[
            1] or W != getattr(self, base_loc_name).shape[2]:
            base_step = down_ratio
            shifts_x = torch.arange(
                0, (W - 1) * base_step + 1,
                base_step,
                dtype=torch.float32,
                device=heatmap.device)
            shifts_y = torch.arange(
                0, (H - 1) * base_step + 1,
                base_step,
                dtype=torch.float32,
                device=heatmap.device)
            shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
            setattr(self, base_loc_name, torch.stack((shift_x, shift_y), dim=0))  # (2, h, w)

        # (batch, h, w, 4)
        pred_boxes = torch.cat((getattr(self, base_loc_name) - pred_wh[:, [0, 1]],
                                getattr(self, base_loc_name) + pred_wh[:, [2, 3]]),
                               dim=1).permute(0, 2, 3, 1)
        boxes = box_target.permute(0, 2, 3, 1)

        mask = wh_weight.view(-1, H, W)
        avg_factor = mask.sum() + 1e-4

        loss_bbox = giou_loss(
            pred_boxes, boxes, mask, avg_factor=avg_factor) * wh_weight_factor

        return loss_cls, loss_bbox

    @force_fp32(apply_to=('pred_hm_b1', 'pred_hm_b2', 'pred_wh_b1', 'pred_wh_b2'))
    def loss(self,
             pred_hm_b1,
             pred_hm_b2,
             pred_wh_b1,
             pred_wh_b2,
             gt_bboxes,
             gt_labels,
             img_metas,
             cfg,
             gt_bboxes_ignore=None):
        h_b1, h_b2, b_b1, b_b2, r_b1, r_b2 = self.ttf_target(
            gt_bboxes, gt_labels, img_metas)

        loss_cls_b1, loss_bbox_b1 = self.loss_single(
            pred_hm_b1, pred_wh_b1, h_b1, b_b1, r_b1,
            self.down_ratio_b1, 'base_loc_b1', self.hm_weight_b1, self.wh_weight_b1,
            4 if self.focal_b2_only else self.focal_loss_beta)
        loss_cls_b2, loss_bbox_b2 = self.loss_single(
            pred_hm_b2, pred_wh_b2, h_b2, b_b2, r_b2,
            self.down_ratio_b2, 'base_loc_b2', self.hm_weight_b2, self.wh_weight_b2,
            self.focal_loss_beta)

        return {'losses/ttfnetv2_loss_hm_b1': loss_cls_b1,
                'losses/ttfnetv2_loss_wh_b1': loss_bbox_b1,
                'losses/ttfnetv2_loss_hm_b2': loss_cls_b2,
                'losses/ttfnetv2_loss_wh_b2': loss_bbox_b2}

    def _topk(self, scores, topk):
        batch, cat, height, width = scores.size()

        # (batch, 80, topk)
        topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), topk)

        topk_inds = topk_inds % (height * width)
        topk_ys = (topk_inds / width).int().float()
        topk_xs = (topk_inds % width).int().float()

        # (batch, topk). select topk from 80*topk
        topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), topk)
        topk_clses = (topk_ind / topk).int()
        topk_ind = topk_ind.unsqueeze(2)
        topk_inds = topk_inds.view(batch, -1,
                                   1).gather(1, topk_ind).view(batch, topk)
        topk_ys = topk_ys.view(batch, -1,
                               1).gather(1, topk_ind).view(batch, topk)
        topk_xs = topk_xs.view(batch, -1,
                               1).gather(1, topk_ind).view(batch, topk)

        return topk_score, topk_inds, topk_clses, topk_ys, topk_xs

    def gaussian_2d(self, shape, sigma_x=1, sigma_y=1):
        m, n = [(ss - 1.) / 2. for ss in shape]
        y, x = np.ogrid[-m:m + 1, -n:n + 1]

        h = np.exp(-(x * x / (2 * sigma_x * sigma_x) + y * y /
                     (2 * sigma_y * sigma_y)))
        h[h < np.finfo(h.dtype).eps * h.max()] = 0
        return h

    def draw_truncate_gaussian(self, heatmap, center, h_radius, w_radius, k=1):
        h, w = 2 * h_radius + 1, 2 * w_radius + 1
        sigma_x = w / 6
        sigma_y = h / 6
        gaussian = self.gaussian_2d((h, w), sigma_x=sigma_x, sigma_y=sigma_y)
        gaussian = heatmap.new_tensor(gaussian)

        x, y = int(center[0]), int(center[1])

        height, width = heatmap.shape[0:2]

        left, right = min(x, w_radius), min(width - x, w_radius + 1)
        top, bottom = min(y, h_radius), min(height - y, h_radius + 1)

        masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
        masked_gaussian = gaussian[h_radius - top:h_radius + bottom, w_radius -
                                                                     left:w_radius + right]
        if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:
            torch.max(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
        return heatmap

    def ttf_target_single_single(self,
                                 heatmap,
                                 box_target,
                                 reg_weight,
                                 fake_heatmap,
                                 boxes_area_topk_log,
                                 gt_boxes,
                                 gt_labels,
                                 boxes_ind,
                                 feat_shape,
                                 down_ratio):
        output_h, output_w = feat_shape
        feat_gt_boxes = gt_boxes / down_ratio
        feat_gt_boxes[:, [0, 2]] = torch.clamp(
            feat_gt_boxes[:, [0, 2]], min=0, max=output_w - 1)
        feat_gt_boxes[:, [1, 3]] = torch.clamp(
            feat_gt_boxes[:, [1, 3]], min=0, max=output_h - 1)
        feat_hs, feat_ws = (feat_gt_boxes[:, 3] - feat_gt_boxes[:, 1],
                            feat_gt_boxes[:, 2] - feat_gt_boxes[:, 0])

        ct_ints = (torch.stack([(gt_boxes[:, 0] + gt_boxes[:, 2]) / 2,
                                (gt_boxes[:, 1] + gt_boxes[:, 3]) / 2],
                               dim=1) / down_ratio).to(torch.int)

        h_radiuses_alpha = (feat_hs / 2. * self.alpha).int()
        w_radiuses_alpha = (feat_ws / 2. * self.alpha).int()
        if self.alpha != self.beta:
            h_radiuses_beta = (feat_hs / 2. * self.beta).int()
            w_radiuses_beta = (feat_ws / 2. * self.beta).int()

        # larger boxes have lower priority than small boxes.
        for k in range(boxes_ind.shape[0]):
            cls_id = gt_labels[k] - 1

            fake_heatmap = fake_heatmap.zero_()
            self.draw_truncate_gaussian(fake_heatmap, ct_ints[k],
                                        h_radiuses_alpha[k].item(),
                                        w_radiuses_alpha[k].item())
            heatmap[cls_id] = torch.max(heatmap[cls_id], fake_heatmap)

            if self.alpha != self.beta:
                fake_heatmap = fake_heatmap.zero_()
                self.draw_truncate_gaussian(fake_heatmap, ct_ints[k],
                                            h_radiuses_beta[k].item(),
                                            w_radiuses_beta[k].item())
            box_target_inds = fake_heatmap > 0

            box_target[:, box_target_inds] = gt_boxes[k][:, None]
            cls_id = 0

            local_heatmap = fake_heatmap[box_target_inds]
            ct_div = local_heatmap.sum()
            local_heatmap *= boxes_area_topk_log[k]
            reg_weight[cls_id, box_target_inds] = local_heatmap / ct_div
        return heatmap, box_target, reg_weight

    def ttf_target_single(self, gt_boxes, gt_labels, feat_shape):
        """

        Args:
            gt_boxes: tensor, tensor <=> img, (num_gt, 4).
            gt_labels: tensor, tensor <=> img, (num_gt,).
            feat_shape: tuple.

        Returns:
            heatmap: tensor, tensor <=> img, (80, h, w).
            box_target: tensor, tensor <=> img, (4, h, w).
            reg_weight: tensor, same as box_target
        """
        output_h_b1, output_w_b1, output_h_b2, output_w_b2 = feat_shape
        heatmap_channel = self.num_fg

        heatmap_b1 = gt_boxes.new_zeros((heatmap_channel, output_h_b1, output_w_b1))
        fake_heatmap_b1 = gt_boxes.new_zeros((output_h_b1, output_w_b1))
        box_target_b1 = gt_boxes.new_ones((4, output_h_b1, output_w_b1)) * -1
        reg_weight_b1 = gt_boxes.new_zeros((1, output_h_b1, output_w_b1))
        heatmap_b2 = gt_boxes.new_zeros((heatmap_channel, output_h_b2, output_w_b2))
        fake_heatmap_b2 = gt_boxes.new_zeros((output_h_b2, output_w_b2))
        box_target_b2 = gt_boxes.new_ones((4, output_h_b2, output_w_b2)) * -1
        reg_weight_b2 = gt_boxes.new_zeros((1, output_h_b2, output_w_b2))

        boxes_areas_log = self.bbox_areas(gt_boxes).log()
        boxes_area_topk_log, boxes_ind = torch.topk(boxes_areas_log,
                                                    boxes_areas_log.size(0))

        gt_boxes = gt_boxes[boxes_ind]
        gt_labels = gt_labels[boxes_ind]

        if self.level_base_area:
            gt_b1_idx = boxes_area_topk_log >= math.log(self.b1_min_length ** 2)
            gt_b2_idx = boxes_area_topk_log <= math.log(self.b2_max_length ** 2)
        else:
            gt_b1_idx = gt_boxes.max(-1)[0] >= self.b1_min_length
            gt_b2_idx = gt_boxes.max(-1)[0] <= self.b2_max_length
        heatmap_b1, box_target_b1, reg_weight_b1 = self.ttf_target_single_single(
            heatmap_b1,
            box_target_b1,
            reg_weight_b1,
            fake_heatmap_b1,
            boxes_area_topk_log[gt_b1_idx],
            gt_boxes[gt_b1_idx],
            gt_labels[gt_b1_idx],
            boxes_ind[gt_b1_idx],
            [output_h_b1, output_w_b1],
            self.down_ratio_b1)

        heatmap_b2, box_target_b2, reg_weight_b2 = self.ttf_target_single_single(
            heatmap_b2,
            box_target_b2,
            reg_weight_b2,
            fake_heatmap_b2,
            boxes_area_topk_log[gt_b2_idx],
            gt_boxes[gt_b2_idx],
            gt_labels[gt_b2_idx],
            boxes_ind[gt_b2_idx],
            [output_h_b2, output_w_b2],
            self.down_ratio_b2)

        return heatmap_b1, heatmap_b2, box_target_b1, box_target_b2, reg_weight_b1, reg_weight_b2

    def ttf_target(self, gt_boxes, gt_labels, img_metas):
        """

        Args:
            gt_boxes: list(tensor). tensor <=> image, (gt_num, 4).
            gt_labels: list(tensor). tensor <=> image, (gt_num,).
            img_metas: list(dict).

        Returns:
            heatmap: tensor, (batch, 80, h, w).
            box_target: tensor, (batch, 4, h, w).
            reg_weight: tensor, (batch, 1, h, w).
        """
        with torch.no_grad():
            feat_shape = (img_metas[0]['pad_shape'][0] // self.down_ratio_b1,
                          img_metas[0]['pad_shape'][1] // self.down_ratio_b1,
                          img_metas[0]['pad_shape'][0] // self.down_ratio_b2,
                          img_metas[0]['pad_shape'][1] // self.down_ratio_b2)

            h_b1, h_b2, b_b1, b_b2, r_b1, r_b2 = multi_apply(
                self.ttf_target_single,
                gt_boxes,
                gt_labels,
                feat_shape=feat_shape)

            h_b1, h_b2, b_b1, b_b2, r_b1, r_b2 = [
                torch.stack(t, dim=0).detach()
                for t in [h_b1, h_b2, b_b1, b_b2, r_b1, r_b2]
            ]

            return h_b1, h_b2, b_b1, b_b2, r_b1, r_b2

    def simple_nms(self, heat, kernel=3, out_heat=None):
        pad = (kernel - 1) // 2
        hmax = nn.functional.max_pool2d(
            heat, (kernel, kernel), stride=1, padding=pad)
        keep = (hmax == heat).float()
        out_heat = heat if out_heat is None else out_heat
        return out_heat * keep

    def bbox_areas(self, bboxes, keep_axis=False):
        x_min, y_min, x_max, y_max = bboxes[:, 0], bboxes[:, 1], \
                                     bboxes[:, 2], bboxes[:, 3]
        areas = (y_max - y_min + 1) * (x_max - x_min + 1)
        if keep_axis:
            return areas[:, None]
        return areas


import alles, sys

import pydub
import loris
import time
import numpy as np
from math import pi
from collections import deque

tests = [
    "/Users/bwhitman/sounds/aps/samples/ADVORCH1/BA LONG FF/BA LGFF#C3.wav",
    "/Users/bwhitman/sounds/aps/samples/ADVORCH2/VIS LONG F/VIS LGF G2.wav",
    "/Users/bwhitman/sounds/aps/samples/AMTMODERNCOMPOSER/DAMP GUITAR/DAMP GTR G2.wav",
    "/Users/bwhitman/sounds/aps/samples/AMTMODERNCOMPOSER/BA HARMONICS/DBL HARM G3.wav",
    "/Users/bwhitman/sounds/aps/samples/AMTMODERNCOMPOSER/BAROQUE MIX/BP C3.wav",
    "/Users/bwhitman/sounds/aps/samples/AMTMODERNCOMPOSER/ALPHA BASS/ALPHA B.-C2.wav"
]

def list_from_py2_iterator(obj, how_many):
    # Oof, the loris object uses some form of iteration that py3 doesn't like.
    ret = []
    it = obj.iterator()
    for i in range(how_many):
        ret.append(it.next())
    return ret

def loris_synth(filename, freq_res=150, analysis_window=100,amp_floor=-30, max_len_s = 10, noise_ratio=1, hop_time=0.04):
    # Pure loris synth for A/B testing
    audio = pydub.AudioSegment.from_file(filename)
    audio = audio[:int(max_len_s*1000.0)]
    y = np.array(audio.get_array_of_samples())
    if audio.channels == 2:
        y =y.reshape((-1, 2))
        y = y[:,1]
    y = np.float64(y) / 2**15
    analyzer = loris.Analyzer(freq_res, analysis_window)
    analyzer.setAmpFloor(amp_floor)
    analyzer.setHopTime(hop_time)
    partials = analyzer.analyze(y,44100)
    bps = 0
    for i in list_from_py2_iterator(partials, len(partials)):
        bps = bps + len(list_from_py2_iterator(i, i.numBreakpoints()))
    print("%d partials %d bps" % (len(partials), bps))
    loris.scaleNoiseRatio(partials, noise_ratio)
    return loris.synthesize(partials,44100)

def sequence(filename, max_len_s = 10, amp_floor=-30, hop_time=0.04, max_oscs=alles.OSCS, freq_res = 10, freq_drift=20, analysis_window = 100):
    # my job: take a file, analyze it, output a sequence + some metadata
    # i do voice stealing to keep maximum partials at once to max_oscs
    # my sequence is an ordered list of partials/oscillators, a list with (ms, osc, freq, amp, bw, phase, time_delta, amp_delta, freq_delta, bw_delta)
    audio = pydub.AudioSegment.from_file(filename)
    audio = audio[:int(max_len_s*1000.0)]
    y = np.array(audio.get_array_of_samples())
    if int(audio.frame_rate) != int(alles.SAMPLE_RATE):
        print("SR mismatch, todo")
        return (None, None)
    if audio.channels == 2:
        y =y.reshape((-1, 2))
        y = y[:,1]
    y = np.float64(y) / 2**15
    metadata = {"filename":filename, "samples":y.shape[0]}

    if(filename.endswith(".wav")):
        import wavdumper # Forked version
        w = wavdumper.Wav(filename)
        w.printInfo() # TODO, call this extractInfo
        try:
            if(w.midiUnityNote>0):
                metadata["midi_note"] = w.midiUnityNote
            if(w.loopstart >= 0 and w.loopend >= 0):
                metadata["sustain_ms"] = int(((w.loopstart + ((w.loopend-w.loopstart)/2.0)) / alles.SAMPLE_RATE) * 1000.0)
        except AttributeError:
            pass # No wav metadata

    # Do the loris analyze
    analyzer = loris.Analyzer(freq_res, analysis_window)
    analyzer.setAmpFloor(amp_floor)
    analyzer.setFreqDrift(freq_drift)
    analyzer.setHopTime(hop_time)
    partials_it = analyzer.analyze(y, audio.frame_rate)
    # build the sequence
    sequence = []
    partials = list_from_py2_iterator(partials_it, partials_it.size())
    partial_count = 0
    for partial_idx, partial in enumerate(partials):
        breakpoints = list_from_py2_iterator(partial, partial.numBreakpoints())
        if(len(breakpoints)>1):
            partial_count = partial_count + 1
            for bp_idx, bp in enumerate(breakpoints):
                phase = -1
                # Last breakpoint?
                if(bp_idx == len(breakpoints)-1): phase = -2
                # First breakpoint
                if(bp_idx == 0):
                    phase = bp.phase() / (2*pi)
                    if(phase < 0): phase = phase + 1
                time_ms = int(bp.time() * 1000.0)
                sequence.append( [time_ms, partial_idx, bp.frequency(), bp.amplitude(), bp.bandwidth(), phase] )

    # Now go and order them and figure out which oscillator gets which partial
    time_ordered = sorted(sequence, key=lambda x:x[0])
    first_time = time_ordered[0][0]
    # Clear the sequence
    sequence = []
    min_q_len = max_oscs
    # Now add in a voice / osc #
    osc_map = {}
    osc_q = deque(range(max_oscs))
    for i,s in enumerate(time_ordered):
        next_idx = -1
        time_delta, amp_delta, freq_delta, bw_delta = (0,0,0,0)
        if(s[5] != -2): # if not the end of a partial
            next_idx = i+1
            while(time_ordered[next_idx][1] != s[1]):
                next_idx = next_idx + 1
            n = time_ordered[next_idx]
            time_delta = n[0] - s[0]
            amp_delta = n[3]/s[3]
            freq_delta = n[2]/s[2]
            if(s[4]>0):
                bw_delta = n[4]/s[4]
            else:
                bw_delta = 0

        s.append(time_delta)
        s.append(amp_delta)
        s.append(freq_delta)
        s.append(bw_delta)

        # Start the partials at 0
        s[0] = s[0] - first_time

        if(s[5]>=0): #new partial
            if(len(osc_q)):
                osc_map[s[1]] = osc_q.popleft()
                # Replace the partial_idx with a osc offset
                s[1] = osc_map[s[1]]
                sequence.append(s)
        else:
            osc = osc_map.get(s[1], None)
            if(osc is not None):
                s[1] = osc_map[s[1]]
                sequence.append(s)
                if(s[5] == -2): # last bp
                    # Put the oscillator back
                    osc_q.appendleft(osc)
        if(len(osc_q) < min_q_len): min_q_len = len(osc_q)
    print("%d partials and %d breakpoints, max oscs used at once was %d" % (partial_count, len(sequence), max_oscs - min_q_len))
    # Fix sustain_ms
    if(metadata.get("sustain_ms", 0) > 0):
        metadata["sustain_ms"] = metadata["sustain_ms"] - first_time
    metadata["oscs_alloc"] = max_oscs-min_q_len
    return (metadata, sequence)

def play(sequence, osc_offset=0, sustain_ms = -1, sustain_len_ms = 0, time_ratio = 1, pitch_ratio = 1, amp_ratio = 1, bw_ratio = 1, round_robin=False):
    # i take a sequence and play it to AMY, just like native AMY will do from a .h file
    my_start_time = alles.millis()

    if(round_robin):
        next_client = 0
        osc_to_client_map = {}
        print("Syncing mesh....")
        clients = len(alles.sync())
        # After a sync, we don't want to immediately spam the mesh, so let's wait 2000ms
        time.sleep(2)
        print("Ready to play among %d speakers" % (clients))

    sustain_offset = 0
    if(sustain_ms > 0):
        if(sustain_ms > sequence[-1][0]):
            print("Moving sustain_ms from %d to %d" % (sustain_ms, sequence[-1][0]-100))
            sustain_ms = sequence[-1][0] - 100

    for i,s in enumerate(sequence):
        # Wait for the item in the sequence to be close, so I don't overflow the synthesizers' state
        while(my_start_time + (s[0] / time_ratio) > (alles.millis() - 500)):
            time.sleep(0.01)

        # Make envelope strings
        bp0 = "%d,%s,0,0" % (s[6] / time_ratio, alles.trunc(s[7]))
        bp1 = "%d,%s,0,0" % (s[6] / time_ratio, alles.trunc(s[8]))
        if(bw_ratio > 0):
            bp2 = "%d,%s,0,0" % (s[6] / time_ratio, alles.trunc(s[9]))
        else:
            bp2 = ""
        if(sustain_ms > 0 and sustain_offset == 0):
            if(s[0]/time_ratio > sustain_ms/time_ratio):
                sustain_offset = sustain_len_ms/time_ratio

        osc = s[1]+osc_offset

        partial_args = {}

        if(round_robin):
            if(osc_to_client_map.get(osc,None) is None):
                osc_to_client_map[osc] = (next_client % clients)
                next_client += 1
            partial_args["client"] = osc_to_client_map[osc]

        partial_args.update({"timestamp":my_start_time + (s[0]/time_ratio + sustain_offset),
            "osc":s[1]+osc_offset,
            "wave":alles.PARTIAL,
            "amp":s[3]*amp_ratio,
            "freq":s[2]*pitch_ratio,
            "feedback":s[4]*bw_ratio,
            "bp0":bp0, "bp1":bp1, "bp2":bp2,
            "bp0_target":alles.TARGET_AMP+alles.TARGET_LINEAR,
            "bp1_target":alles.TARGET_FREQ+alles.TARGET_LINEAR,
            "bp2_target":alles.TARGET_FEEDBACK+alles.TARGET_LINEAR})

        if(s[5]==-2): #end, add note off
            alles.send(**partial_args, vel=0)
        elif(s[5]==-1): # continue
            alles.send(**partial_args)
        else: #start, add phase and note on
            alles.send(**partial_args, vel=s[3]*amp_ratio, phase=s[5])

    return sequence[-1][0]/time_ratio

#In [6]: partials.generate_partials_header(fns,amp_floor=-40,analysis_window=40,freq_drift=5,hop_time=0.04,freq_res=5)
def generate_partials_header(filenames, **kwargs):
    # given a list of filenames, output a partials.h
    out = open("main/amy/partials.h", "w")
    out.write("// Automatically generated by partials.generate_partials_header()\n#ifndef __PARTIALS_H\n#define __PARTIALS_H\n#define PARTIALS_PATCHES %d\n" % (len(filenames)))
    all_partials = []
    for f in filenames:
        m, s = sequence(f, **kwargs)
        if(m is not None):
            all_partials.append((m ,s))
    out.write("const partial_breakpoint_map_t partial_breakpoint_map[%d] = {\n" % (len(all_partials)))
    out.write("\t// offset, length, midi_note, sustain_ms, oscs_alloc\n")
    start = 0
    for p in all_partials:
        out.write("\t{ %d, %d, %d, %d, %d }, /* %s */ \n" % (start, len(p[1]), p[0].get("midi_note", 0), p[0].get("sustain_ms", 0),  p[0]["oscs_alloc"], p[0]["filename"]))
        start = start + len(p[1])
    out.write("};\n");
    out.write("const partial_breakpoint_t partial_breakpoints[%d] = {\n" % (start))
    out.write("\t// ms_offset, osc, freq, amp, bw, phase, ms_delta, amp_delta, freq_delta, bw_delta\n")
    for p in all_partials:
        for s in p[1]:
            out.write("\t { %d, %d, %f, %f, %f, %f, %d, %f, %f, %f }, \n" % tuple(s))
    out.write("};\n")
    out.write("#endif // __PARTIALS_H\n")
    out.close()

# OK defaults here
def test(   filename="/Users/bwhitman/sounds/billboard/0157/0157.mp4", \
                    max_len_s=60, \
                    freq_res = 10, \
                    analysis_window = 100, \
                    time_ratio = 1, \
                    max_oscs = 40, \
                    bw_ratio = 0,\
                    amp_ratio = 1, \
                    pitch_ratio = 1, \
                    amp_floor = -40, \
                    hop_time = 0.04, \
                    sustain_len_ms = 0, \
                    **kwargs):
    import sounddevice as sd
    alles.stop()
    alles.start(immediate=False)
    m,s = sequence(filename, max_len_s = max_len_s, freq_res = freq_res, analysis_window = analysis_window, amp_floor=amp_floor, hop_time=hop_time, max_oscs=max_oscs)
    ms = play(s, sustain_ms = m.get("sustain_ms", -1), time_ratio=time_ratio, pitch_ratio=pitch_ratio, amp_ratio=amp_ratio, bw_ratio = bw_ratio, sustain_len_ms = sustain_len_ms)
    sd.play(alles.render(ms/1000.0))


from __future__ import division

import os
import argparse
import tqdm
import random
import numpy as np

from PIL import Image

import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.autograd import Variable
from pytorchyolo.models import load_model
from pytorchyolo.utils.utils import load_classes, rescale_boxes, non_max_suppression, print_environment_info
from pytorchyolo.utils.datasets import ImageFolder
from pytorchyolo.utils.transforms import Resize, DEFAULT_TRANSFORMS
from common_utils import vis
from common_utils import AlgBase
import yaml
import cv2

class Alg(AlgBase):
    def __init__(self):
        self.cfg_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "cfg.yaml")
        self.cfg_info = {}
        self.model_name = None
        self.model = None
        self.device = "cpu"
        self.ignore_keys = []
        self.load_cfg()

    def create_model(self, model_name='yolov3_tiny', dev='cpu'):
        if model_name not in self.cfg_info.keys():
            print('error: unknown model name ', model_name, 'create failed')
            return
        self.device = dev
        self.model_name = model_name

        model_cfg = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'config/%s'%(self.cfg_info[model_name]['normal']['cfg_file']))
        model_weight = os.path.join(os.path.dirname(os.path.realpath(__file__)), '%s'%(self.cfg_info[model_name]['normal']['weight']))
        self.model = load_model(model_cfg, self.device, model_weight)
        self.model.eval()

    def inference(self, img_array):
        map_result = {'type':'img'}
        img_tensor = transforms.Compose([
            DEFAULT_TRANSFORMS,
            Resize(int(self.cfg_info[self.model_name]['normal']['infer_size']))])(
                (img_array, np.zeros((1, 5))))[0].unsqueeze(0)
        if self.device == "cuda":
            img_tensor = img_tensor.cuda()

        with torch.no_grad():
            detections = self.model(img_tensor)
            detections = non_max_suppression(detections, float(self.cfg_info[self.model_name]['normal']['infer_conf']), float(self.cfg_info[self.model_name]['normal']['nms_thre']))
            detections = rescale_boxes(detections[0], int(self.cfg_info[self.model_name]['normal']['infer_size']), img_array.shape[:2])

        valid_pred = detections.cpu()
        boxes = valid_pred[:,0:4]
        cls = valid_pred[:, 5]
        scores = valid_pred[:, 4]
        vis(img_array, boxes, scores, cls, conf=0.0)
        map_result['result'] = img_array
        return map_result


# Make random single-end reads from restriction digest of genome. Currently must be palindromic.
# Will need to update if using non-palindromic cutter.

from optparse import OptionParser
from Bio import SeqIO
import re
import gzip
import numpy as np

def parse_options():
	parser = OptionParser()
	parser.add_option("-g", "--genome_file", dest="genome_file",
					  help="Genome file, fasta format", metavar="GENOMEFILE")
	parser.add_option("-o", "--outfile", dest="outfile",
					  help="Path for output file", metavar="OUTFILE")
	parser.add_option("-n", "--reads_per_cut", dest="reads_per_cut", default=200,
					  help="Number of reads for each restriction site (cut)", metavar="N")
	parser.add_option("-c", "--cut_seq", dest="cut_seq", default='AATT',
					  help="Palindromic restriction site (default AATT)", metavar="Cutseq")

	(options, args) = parser.parse_args()
	return options

options = parse_options()
genome_file = options.genome_file
outfile = gzip.open(options.outfile + '.gz','wb')
reads_per_cut = int(options.reads_per_cut)
restriction_site = options.cut_seq

def revcomp(input):
    output = ''
    for letter in input:
        letter = letter.upper()

        if letter == 'A':
            output += 'T'
        elif letter == 'T':
            output += 'A'
        elif letter == 'G':
            output += 'C'
        else:
            output += 'G'

    return(output[::-1])

seq_id = 1

for record in SeqIO.parse(genome_file, "fasta"):
    chr_ = re.sub('chr', '', record.id)
    print(chr_)
    chr_seq = str(record.seq).upper()
    cut_sites = [m.start() for m in re.finditer(restriction_site, chr_seq)]
    for site in cut_sites:
        # Draw 5 just because.
        for n in range(0, reads_per_cut):
            shear_length = np.max([np.random.poisson(150,1)[0], 100])
            direction = np.random.choice([-1,1])
            if (direction == -1):
                start = site - shear_length
                end = site
                seq = chr_seq[start:end]
            elif (direction == 1):
                start = site
                end = start + shear_length
                seq = revcomp(chr_seq[start:end])
            if (len(seq) >= 100):
                seq = seq[:100]
                quals = '~' * 100
                line_towrite = '@' + str(seq_id) + '\n' + seq + '\n+\n' +  quals + '\n'
                outfile.write(line_towrite.encode())

            seq_id += 1

outfile.close()


# This file is part of the pyBinSim project.
#
# Copyright (c) 2017 A. Neidhardt, F. Klein, N. Knoop, T. Köllmer
#
# 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 logging
import multiprocessing

import numpy as np
import soundfile as sf

from pybinsim.pose import Pose
from pybinsim.utility import total_size

nThreads = multiprocessing.cpu_count()

class Filter(object):

    def __init__(self, inputfilter, irBlocks, block_size):

        self.IR_left_blocked = np.reshape(inputfilter[:, 0], (irBlocks, block_size))
        self.IR_right_blocked = np.reshape(inputfilter[:, 1], (irBlocks, block_size))

    def getFilter(self):
        return self.IR_left_blocked, self.IR_right_blocked

class FilterStorage(object):
    """ Class for storing all filters mentioned in the filter list """

    def __init__(self, irSize, block_size, filter_list_name):

        self.log = logging.getLogger("pybinsim.FilterStorage")
        self.log.info("FilterStorage: init")

        self.ir_size = irSize
        self.ir_blocks = irSize // block_size
        self.block_size = block_size
        self.default_filter = Filter(np.zeros((self.ir_size, 2), dtype='float32'),self.ir_blocks,self.block_size)

        self.filter_list_path = filter_list_name
        self.filter_list = open(self.filter_list_path, 'r')

        self.headphone_filter = None

        # format: [key,{filter}]
        self.filter_dict = {}

        # Start to load filters
        self.load_filters()

    def parse_filter_list(self):
        """
        Generator for filter list lines

        Lines are assumed to have a format like
        0 0 40 1 1 0 brirWav_APA/Ref_A01_1_040.wav

        The headphone filter starts with HPFILTER instead of the positions.

        Lines can be commented with a '#' as first character.

        :return: Iterator of (Pose, filter-path) tuples
        """

        for line in self.filter_list:

            # comment out lines in the list with a '#'
            if line.startswith('#'):
                continue

            line_content = line.split()
            filter_path = line_content[-1]

            if line.startswith('HPFILTER'):
                self.log.info("Loading headphone filter: {}".format(filter_path))
                self.headphone_filter = Filter(self.load_filter(filter_path),self.ir_blocks,self.block_size)
                continue

            filter_value_list = tuple(line_content[0:-1])

            pose = Pose.from_filterValueList(filter_value_list)

            yield pose, filter_path

    def load_filters(self):
        """
        Load filters from files

        :return: None
        """

        self.log.info("Start loading filters...")

        for i, (pose, filter_path) in enumerate(self.parse_filter_list()):
            self.log.debug('Loading {}'.format(filter_path))

            current_filter = Filter(self.load_filter(filter_path), self.ir_blocks, self.block_size)

            # create key and store in dict.
            key = pose.create_key()
            self.filter_dict.update({key: current_filter})

        self.log.info("Finished loading filters.")
        #self.log.info("filter_dict size: {}MiB".format(total_size(self.filter_dict) // 1024 // 1024))

    def get_filter(self, pose):
        """
        Searches in the dict if key is available and return corresponding filter
        When no filter is found, defaultFilter is returned which results in silence

        :param pose
        :return: corresponding filter for pose
        """

        key = pose.create_key()

        if key in self.filter_dict:
            self.log.info('Filter found: key: {}'.format(key))
            return self.filter_dict.get(key)
        else:
            self.log.warning('Filter not found: key: {}'.format(key))
            return self.default_filter

    def close(self):
        self.log.info('FilterStorage: close()')
        # TODO: do something in here?

    def get_headphone_filter(self):
        if self.headphone_filter is None:
            raise RuntimeError("Headphone filter not loaded")

        return self.headphone_filter

    def load_filter(self, filter_path):

        current_filter, fs = sf.read(filter_path, dtype='float32')

        filter_size = np.shape(current_filter)

        # Fill filter with zeros if to short
        if filter_size[0] < self.ir_size:
            self.log.warning('Filter to short: Fill up with zeros')
            current_filter = np.concatenate((current_filter, np.zeros((self.ir_size - filter_size[0], 2))), 0)

        return current_filter


from copy import deepcopy
from shutil import copyfile
from sys_simulator.q_learning.environments.completeEnvironment12 import CompleteEnvironment12
from time import time

import numpy as np
import torch
from torch.utils.tensorboard import SummaryWriter

from sys_simulator.channels import BANChannel, UrbanMacroNLOSWinnerChannel
from sys_simulator.ddpg.agent import SurrogateAgent, SysSimAgent, SysSimAgentWriter
from sys_simulator.ddpg.framework import Framework
from sys_simulator.devices.devices import db_to_power
import sys_simulator.general as gen
from sys_simulator.general import power_to_db, print_evaluating, print_stuff_ddpg
from sys_simulator.general.ou_noise import SysSimOUNoise
from sys_simulator.parameters.parameters import EnvironmentParameters
from sys_simulator.plots import plot_env_states
from sys_simulator.q_learning.environments.completeEnvironment11 import (
    CompleteEnvironment11,
)

# Uses CompleteEnvironment11
# Centralized Learning-Centralized Execution
# Trains a central agent for a fixed amount of devices.
# There are different channels to the BS and to the devices.
# The script generates a single scenario and tries to find
# the optimal solution for it.
# Single episodes convergence. The states are in linear scale.
n_mues = 1  # number of mues
n_d2d = 2  # number of d2d pairs
n_rb = n_mues   # number of RBs
carrier_frequency = 2.4  # carrier frequency in GHz
bs_radius = 500  # bs radius in m
rb_bandwidth = 180*1e3  # rb bandwidth in Hz
d2d_pair_distance = 50  # d2d pair distance in m
device_height = 1.5  # mobile devices height in m
bs_height = 25  # BS antenna height in m
p_max = 40  # max tx power in dBm
noise_power = -116  # noise power per RB in dBm
bs_gain = 17    # macro bs antenna gain in dBi
user_gain = 4   # user antenna gain in dBi
sinr_threshold_train = 6  # mue sinr threshold in dB for training
mue_margin = 200  # mue margin in dB
# conversions from dBm to dB
p_max = p_max - 30
noise_power = noise_power - 30
# env parameters
CHANNEL_RND = True
C = 8  # C constant for the improved reward function
ENVIRONMENT_MEMORY = 2
MAX_NUMBER_OF_AGENTS = 3
REWARD_PENALTY = 1.5
N_STATES_BINS = 100
# q-learning parameters
# training
ALGO_NAME = 'ddpg'
REWARD_FUNCTION = 'classic'
MAX_STEPS = 6000
# MAX_STEPS = 1000
STEPS_PER_EPISODE = MAX_STEPS
EVAL_STEPS_PER_EPISODE = 50
REPLAY_INITIAL = int(1E3)
EVAL_NUM_EPISODES = 5
REPLAY_MEMORY_SIZE = int(1E4)
ACTOR_LEARNING_RATE = 5E-5
CRITIC_LEARNING_RATE = 5E-4
HIDDEN_SIZE = 256
N_HIDDEN_LAYERS = 3
BATCH_SIZE = 128
GAMMA = .99
SOFT_TAU = 1E-2
ALPHA = .6
BETA = .4
EXPLORATION = 'ou'
REPLAY_MEMORY_TYPE = 'standard'
PRIO_BETA_ITS = int(.8*(MAX_STEPS - REPLAY_INITIAL))
PRINT_EVERY = int(MAX_STEPS/100)
EVAL_EVERY = int(MAX_STEPS / 10)
# EVAL_EVERY = int(MAX_STEPS / 1)
OU_DECAY_PERIOD = 100000
OU_MU = 0.0
OU_THETA = .15
OU_MAX_SIGMA = .3
OU_MIN_SIGMA = .3

torch_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
env_params = EnvironmentParameters(
    rb_bandwidth, d2d_pair_distance, p_max, noise_power,
    bs_gain, user_gain, sinr_threshold_train,
    n_mues, n_d2d, n_rb, bs_radius, c_param=C, mue_margin=mue_margin
)
channel_to_devices = BANChannel(rnd=CHANNEL_RND)
channel_to_bs = UrbanMacroNLOSWinnerChannel(
    rnd=CHANNEL_RND, f_c=carrier_frequency, h_bs=bs_height, h_ms=device_height
)
ref_env = CompleteEnvironment11(
    env_params,
    channel_to_bs,
    channel_to_devices,
    reward_penalty=REWARD_PENALTY,
    memory=ENVIRONMENT_MEMORY,
    bs_height=bs_height,
    reward_function=REWARD_FUNCTION,
    memories_capacity=int(1e3)
)
a_min = 0 + 1e-9
a_max = db_to_power(p_max - 10)
action_size = MAX_NUMBER_OF_AGENTS
env_state_size = MAX_NUMBER_OF_AGENTS * ref_env.state_size()

framework = Framework(
    REPLAY_MEMORY_TYPE,
    REPLAY_MEMORY_SIZE,
    REPLAY_INITIAL,
    env_state_size,
    action_size,
    HIDDEN_SIZE,
    N_HIDDEN_LAYERS,
    ACTOR_LEARNING_RATE,
    CRITIC_LEARNING_RATE,
    BATCH_SIZE,
    GAMMA,
    SOFT_TAU,
    torch_device,
    alpha=ALPHA,
    beta=BETA,
    beta_its=PRIO_BETA_ITS
)

ou_noise = SysSimOUNoise(
    action_size,
    a_min, a_max,
    OU_MU, OU_THETA,
    OU_MAX_SIGMA,
    OU_MIN_SIGMA,
    OU_DECAY_PERIOD
)

central_agent = SysSimAgentWriter(a_min, a_max, EXPLORATION, torch_device)
central_agent_test = SysSimAgent(a_min, a_max, EXPLORATION, torch_device)
surr_agents = [SurrogateAgent() for _ in range(MAX_NUMBER_OF_AGENTS)]

def train(env: CompleteEnvironment11, start: float, writer: SummaryWriter):
    actor_losses_bag = list()
    critic_losses_bag = list()
    mue_spectral_eff_bag = list()
    d2d_spectral_eff_bag = list()
    rewards_bag = list()
    mue_avail_bag = list()
    step = 0
    collected_states = list()
    while step < MAX_STEPS:
        obs_aux, _, _, _ = env.step(surr_agents)
        obs = np.concatenate(obs_aux, axis=1)
        total_reward = 0.0
        done = False
        i = 0
        while not done and i < STEPS_PER_EPISODE:
            actions = central_agent.act(obs, framework,
                                        writer, step, True, step=i, ou=ou_noise)
            for j, agent in enumerate(surr_agents):
                agent.set_action(actions[0][j])
            next_obs_aux, rewards, done, _ = env.step(surr_agents)
            collected_states += next_obs_aux
            total_reward = np.sum(rewards)
            next_obs = np.concatenate(next_obs_aux, axis=1)
            framework.replay_memory.push(obs, actions, total_reward, next_obs,
                                         done)
            actor_loss, critic_loss = framework.learn()
            obs = next_obs
            i += 1
            step += 1
            writer.add_scalar('Actor Losses', actor_loss, step)
            writer.add_scalar('Critic Losses', critic_loss, step)
            aux_actions = actions.cpu().numpy()
            writer.add_scalar('Average Actions', aux_actions.mean(), step)
            # writer.add_scalar('Aggregated Actions', aux_actions.sum(), step)
            actor_losses_bag.append(actor_loss)
            critic_losses_bag.append(critic_loss)
            framework.actor.train()
            framework.critic.train()
            if step % PRINT_EVERY == 0:
                now = (time() - start) / 60
                print_stuff_ddpg(step, now, MAX_STEPS, REPLAY_MEMORY_TYPE)
            if step % EVAL_EVERY != 0:
                continue
            # testing
            t_bags = test(env, framework)
            t_rewards = t_bags['rewards']
            t_mue_spectral_effs = t_bags['mue_spectral_effs']
            t_d2d_spectral_effs = t_bags['d2d_spectral_effs']
            t_availability = t_bags['mue_availability']
            # mue spectral eff
            mue_spectral_eff_bag.append(t_mue_spectral_effs)
            # average d2d spectral eff
            d2d_spectral_eff_bag.append(t_d2d_spectral_effs)
            rewards_bag.append(t_rewards)
            mue_avail_bag.append(t_availability)
            # write metrics
            writer.add_scalar('Average MUE Spectral efficiencies',
                              np.mean(t_mue_spectral_effs), step)
            writer.add_scalar('Average D2D Spectral Efficiencies',
                              np.mean(t_d2d_spectral_effs), step)
            writer.add_scalar('Aggregated Rewards', np.sum(t_rewards), step)
            writer.add_scalar('Average MUE Availability',
                              np.mean(t_availability), step)
    all_bags = {
        'actor_losses': actor_losses_bag,
        'critic_losses': critic_losses_bag,
        'mue_spectral_effs': mue_spectral_eff_bag,
        'd2d_spectral_effs': d2d_spectral_eff_bag,
        'collected_states': collected_states
    }
    return all_bags

def test(env: CompleteEnvironment11, framework: Framework):
    framework.actor.eval()
    framework.critic.eval()
    mue_availability = []
    mue_spectral_effs = []
    d2d_spectral_effs = []
    rewards_bag = []
    for _ in range(EVAL_NUM_EPISODES):
        obs_aux, _, _, _ = env.step(surr_agents)
        obs = np.concatenate(obs_aux, axis=1)
        i = 0
        done = False
        ep_availability = []
        ep_rewards = []
        ep_mue_speffs = []
        ep_d2d_speffs = []
        while not done and i < EVAL_STEPS_PER_EPISODE:
            actions = central_agent_test.act(obs, framework, False)
            for j, agent in enumerate(surr_agents):
                agent.set_action(actions[0][j])
            next_obs_aux, rewards, done, _ = env.step(surr_agents)
            next_obs = np.concatenate(next_obs_aux, axis=1)
            obs = next_obs
            ep_availability.append(env.mue.sinr > env.params.sinr_threshold)
            ep_rewards.append(np.sum(rewards))
            ep_mue_speffs.append(env.mue_spectral_eff)
            ep_d2d_speffs.append(env.d2d_spectral_eff)
            i += 1
        rewards_bag.append(np.sum(ep_rewards))
        mue_spectral_effs.append(np.mean(ep_mue_speffs))
        d2d_spectral_effs.append(np.mean(ep_d2d_speffs))
        mue_availability.append(np.mean(ep_availability))
    all_bags = {
        'rewards': rewards_bag,
        'mue_spectral_effs': mue_spectral_effs,
        'd2d_spectral_effs': d2d_spectral_effs,
        'mue_availability': mue_availability
    }
    return all_bags

def evaluate(start: float, writer: SummaryWriter, env: CompleteEnvironment11):
    step = 0
    while step < MAX_STEPS:
        obs, _, _, _ = env.step(surr_agents)
        now = (time() - start) / 60
        print_evaluating(step, now, MAX_STEPS)
        done = False
        i = 0
        while not done and i < STEPS_PER_EPISODE:
            actions = central_agent.act(obs, framework,
                                        writer, step, False)
            db_actions = power_to_db(actions.cpu().numpy())
            for j, agent in enumerate(surr_agents):
                agent.set_action(db_actions[j].item())
            next_obs, reward, done, _ = env.step(surr_agents)
            framework.replay_memory.push(obs, actions, reward, next_obs,
                                         done)
            obs = next_obs
            i += 1
            step += 1
            writer.add_scalar('Rewards', reward, step)
            sinrs = [a.d2d_tx.sinr for a in surr_agents]
            sinrs = {f'device {i}': s for i, s in enumerate(sinrs)}
            writer.add_scalars('SINRs', sinrs, step)
            writer.add_scalars(
                'Transmission powers',
                {f'device {i}': a for i, a in enumerate(db_actions)},
                step
            )
            mue_success = int(env.mue.sinr > env.params.sinr_threshold)
            writer.add_scalar('MUE success', mue_success, step)

def run():
    # make data dir
    filename = gen.path_leaf(__file__)
    filename = filename.split('.')[0]
    dir_path = f'data/{ALGO_NAME}/{filename}'
    data_path, _ = gen.make_dir_timestamp(dir_path)
    writer = SummaryWriter(f'{data_path}/tensorboard')
    start = time()
    # set environment up
    env = deepcopy(ref_env)
    env.build_scenario(surr_agents)
    train_bags = train(env, start, writer)
    states = train_bags['collected_states']
    plot_env_states(states, N_STATES_BINS, f'{data_path}/env_states.png')
    writer.close()
    test_bags = test(env, framework)
    # save stuff
    now = (time() - start) / 60
    data_file_path = f'{data_path}/log.pickle'
    data = {
        'train_bags': train_bags,
        'test_bags': test_bags,
        'elapsed_time': now,
        'eval_every': EVAL_EVERY,
    }
    gen.save_with_pickle(data, data_file_path)
    copyfile(__file__, f'{data_path}/{filename}.py')
    print(f'done. Elapsed time: {now} minutes.')

if __name__ == '__main__':
    run()