import multiprocessing
import sys
import time
from datetime import datetime

import scipy
import pickle
import os

from tqdm.autonotebook import tqdm
import numpy as np
import torch

from shapefit.deform.src.deformation.optimization import LBFGS_P
from shapefit.deform.src.deformation.utils import remove_degeneracies, plot_info
from shapefit.deform.src.deformation.transformation import edges_deformation_from_vertices, deformation_matrix
from shapefit.deform.src.deformation.preproc import compute_bitriangles

class Deformer:
    def __init__(self,
                 mesh,
                 target_vertices_transforms,
                 sigma=0.5,
                 voxel_centers_nn=None,
                 surface_samples_nn=None,
                 voxel_centers_p2p=None,
                 surface_samples_p2p=None,
                 unlabeled_points=None,
                 sharp_edges=None,
                 kernel='gauss_kernel',
                 mapping='nn',
                 cuda=None,
                 device_mode='cpu',
                 deg_thr=1e-3,
                 ampl_factor=1,
                 verbose=1):
        """ Class for performing deformations on the set of meshes, labeled with different classes

        Parameters:
            parts (list<trimesh.base.Trimesh>): list of trimesh meshes
            initial_vertices_transforms (list<torch.Tensor>): list of transform matrices given by AlignmentHeatmap [n_parts, (n_vertices, 4, 4)]
            starting_vertices_transforms (list<torch.Tensor>): list of transform matrices given by user-defined deformations [n_parts, (n_vertices, 4, 4)]

        """
        self.cuda = cuda
        self.device_mode = device_mode

        if self.device_mode == 'gpu':
            self.device = [self.cuda[0], self.cuda[0], self.cuda[0], self.cuda[0]] if len(self.cuda) == 1 else \
                    [self.cuda[0], self.cuda[1], self.cuda[1], self.cuda[1]]
            if len(self.cuda) == 4:
                self.device = [self.cuda[0], self.cuda[1], self.cuda[2], self.cuda[3]]
        else:
            self.device = ['cpu', 'cpu', 'cpu', 'cpu']

        self.mesh = mesh
        self.alpha_0 = 1.0
        self.alpha_reg = 1.0
        self.alpha_data = 1.0
        self.verbose = verbose
        self.deg_thr = deg_thr
        self.ampl_factor = ampl_factor

        if verbose == 1:
            print('initializing parts descriptions...')
        self.num_total_vertices = 0
        self.target_vertices = []
        # list of parameters to optimize
        self.segmentation_soft_indicators = []
        self.voxel_centers_nn = voxel_centers_nn
        self.voxel_centers_p2p = voxel_centers_p2p
        self.kernel = kernel
        self.mapping = mapping
        self.surface_samples_nn = surface_samples_nn
        self.surface_samples_p2p = surface_samples_p2p
        self.unlabeled_points = unlabeled_points
        self.precond = None
        self.sharp_edges = sharp_edges
        self.vertices_approximation = None

        self.num_total_vertices += len(mesh.vertices)

        self.mesh_unique_edges = np.array(mesh.edges_unique)
        self.mesh_unique_faces, _ = np.unique(np.sort(mesh.faces, axis=1), axis=0, return_index=True)

        self.torch_vertices = torch.DoubleTensor(mesh.vertices)
        self.torch_unique_edges = torch.IntTensor(self.mesh_unique_edges)
        self.torch_unique_faces = torch.IntTensor(self.mesh_unique_faces)
        self.target_vertices_transforms = target_vertices_transforms.double()
        self.torch_vertices_4d = torch.DoubleTensor(np.hstack([mesh.vertices, np.ones(len(mesh.vertices))[:, None]]))

        torch_tmp_vertices = torch.zeros_like(self.torch_vertices_4d).double()
        for j, transform in enumerate(self.target_vertices_transforms):
            torch_tmp_vertices[j] = torch.mv(transform, self.torch_vertices_4d[j]).double()
        self.target_vertices = torch_tmp_vertices

        if not mesh.is_watertight:
            raise ValueError('Mesh should be watertight')

        if verbose == 1:
            print('computing bitriangles maps...')
        self.bitriangles_map = torch.IntTensor(compute_bitriangles(self.mesh_unique_faces, self.mesh_unique_edges))

        if verbose == 1:
            print('computing faces-to-edges maps...')
        # list<np.array>[n_parts, (n_faces, 3)]
        self.faces_to_edges_map = np.unique(np.sort(mesh.faces_unique_edges, axis=1), axis=0)

        if verbose == 1:
            print('computing adjacent edges for each face...')
        # list<np.array>[n_parts, (3*n_faces, 2)]
        tmp_adjacent_edges = torch.zeros((len(self.faces_to_edges_map) * 3, 2), dtype=torch.int)
        for j, face in enumerate(self.faces_to_edges_map):
            tmp_adjacent_edges[3*j] = torch.IntTensor([face[0], face[1]])
            tmp_adjacent_edges[3*j+1] = torch.IntTensor([face[0], face[2]])
            tmp_adjacent_edges[3*j+2] = torch.IntTensor([face[1], face[2]])
        self.adjacent_edges = tmp_adjacent_edges.long()

        (bitriangles_map_updated,
         torch_vertices_4d_updated,
         target_vertices_updated,
         updated_edges,
         n_vertices,
         n_vertices_old,
         n_edges,
         n_edges_old) = remove_degeneracies(self.torch_vertices_4d,
                                            self.target_vertices,
                                            self.torch_unique_edges,
                                            self.bitriangles_map,
                                            self.deg_thr,
                                            self.ampl_factor)

        self.bitriangles_map_updated = bitriangles_map_updated
        self.torch_vertices_4d_updated = torch_vertices_4d_updated
        self.target_vertices_updated = target_vertices_updated
        self.updated_edges = updated_edges
        self.n_vertices = n_vertices
        self.n_vertices_old = n_vertices_old
        self.n_edges = n_edges
        self.n_edges_old = n_edges_old

        if verbose == 1:
            print('computing edges deformations...')
        # list<torch.Tensor>[n_parts, (n_edges, 4, 4)]
        self.target_edges_deformations = edges_deformation_from_vertices(self.torch_vertices_4d_updated,
                                                                         self.torch_unique_edges,
                                                                         self.target_vertices_updated,
                                                                         self.bitriangles_map_updated)

        if (self.voxel_centers_nn is not None) or (self.voxel_centers_p2p is not None):
            if verbose == 1:
                print('constructing soft indicators...')

            def indicator_factory(a):
                a = torch.Tensor(a).double()
                if self.cuda != None:
                    a = a.to(self.device[0])

                def aux_function(x):
                    kernel_value, count_close_vertices = 0, 0
                    if self.kernel == 'ep_kernel':
                        if self.mapping == 'p2p':
                            distances = torch.sqrt(torch.sum((a[:, None, :] - x[None, ...]) ** 2, dim=2))
                            close_points = (distances < sigma).double()
                            close_points_count = (distances < 0.02).double()
                            count_close_vertices = torch.sum(close_points_count, dim=1)
                            filter_count_close_points = (count_close_vertices <= 0).double()[:, None, None]
                            kernel_value = torch.sum(filter_count_close_points * close_points[..., None] * (a[:, None, :] - x[None, ...]) ** 2, dim=2)
                        elif self.mapping == 'nn':
                            kernel_value = torch.sum((a - x) ** 2, dim=1)

                        if self.mapping != 'p2p':
                            return kernel_value
                        else:
                            return kernel_value, count_close_vertices
                return aux_function

        if self.mapping == 'p2p':
            self.segmentation_soft_indicators_p2p = []
            for i in range(len(self.voxel_centers_p2p)):
                if len(self.voxel_centers_p2p[i]) == 0:
                    self.segmentation_soft_indicators_p2p += [0]
                else:
                    self.segmentation_soft_indicators_p2p += [indicator_factory(self.voxel_centers_p2p[i])]
        else:
            self.segmentation_soft_indicators_nn = indicator_factory(self.voxel_centers_nn)

        self.edges_deformations = []
        if verbose == 1:
            print('initialization done')

        # timing
        self.deformations_time = 0
        self.cost_time = 0
        self.backward_time = 0

    def save_inv_matrix(self, path, alpha_0=1, alpha_reg=1, alpha_sharp=1):

        os.makedirs(path, exist_ok=True)

        deformation_dict = deformation_matrix(self.target_vertices_updated,
                                              self.torch_vertices_4d_updated,
                                              self.torch_unique_edges,
                                              self.bitriangles_map_updated,
                                              self.target_edges_deformations,
                                              self.faces_to_edges_map,
                                              self.torch_unique_faces,
                                              self.sharp_edges,
                                              alpha_0,
                                              alpha_reg,
                                              alpha_sharp)

        self.A_matrices = deformation_dict['A_matrices']

        msize = self.A_matrices[0].shape[0]
        A_1 = self.A_matrices[0][:msize//3, :msize//3].tocsc()
        A_1_inv = scipy.sparse.linalg.inv(A_1)

        with open(os.path.join(path, 'A_1_inv.pkl'), 'wb') as f:
            pickle.dump(A_1_inv, f)

    def solve_data(self,
                   iterations,
                   alpha_0=1,
                   alpha_reg=1,
                   alpha_sharp=1,
                   alpha_data_nn=1,
                   alpha_data_p2p=1,
                   alpha_quad=1,
                   lr=1e-0,
                   print_freq=5,
                   use_precond=False,
                   hessian_cpu=True,
                   load_hessian_path=None,
                   plot=True):

        deformation_dict = deformation_matrix(self.target_vertices_updated,
                                              self.torch_vertices_4d_updated,
                                              self.torch_unique_edges,
                                              self.bitriangles_map_updated,
                                              self.target_edges_deformations,
                                              self.faces_to_edges_map,
                                              self.torch_unique_faces,
                                              self.sharp_edges,
                                              alpha_0,
                                              alpha_reg,
                                              alpha_sharp)

        A_matrices = deformation_dict['A_matrices']
        b_vectors = deformation_dict['b_vectors']
        A_0_sparse = deformation_dict['A_0_sparse']
        A_2_sparse = deformation_dict['A_2_sparse']
        A_3_sparse = deformation_dict['A_3_sparse']
        affine_components = deformation_dict['affine_components']
        translation_components = deformation_dict['translation_components']
        A_0_sparse_sqrt = deformation_dict['A_0_sparse_sqrt']
        D_e_sparse = deformation_dict['D_e_sparse']
        A_2_sparse_sqrt = deformation_dict['A_2_sparse_sqrt']

        vertices = self.torch_vertices_4d_updated[:, :3]
        vertices_concat = torch.cat([vertices[:, 0],
                                     vertices[:, 1],
                                     vertices[:, 2]]).double()
        vertices_target = self.target_vertices_updated[:, :3]

        A_0_sparse /= self.n_vertices
        b_vectors[0] /= self.n_vertices
        A_2_sparse /= self.n_vertices
        A_3_sparse /= self.n_vertices

        A_0_sparse_sqrt /= np.sqrt(self.n_vertices)
        D_e_sparse /= np.sqrt(self.n_vertices)
        A_2_sparse_sqrt /= np.sqrt(self.n_vertices)

        try:
            lin_grad = alpha_0 * A_0_sparse @ vertices_concat.numpy() - alpha_0 * b_vectors[0].toarray().reshape(-1) + alpha_reg * A_2_sparse @ vertices_concat.numpy() + alpha_sharp * A_3_sparse @ vertices_concat.numpy()
            lin_grad_safe = alpha_0 * (A_0_sparse_sqrt.T) @ (A_0_sparse_sqrt @ vertices_concat.numpy()) + \
                            alpha_0 * (D_e_sparse.T) @ (D_e_sparse @ vertices_concat.numpy()) - \
                            alpha_0 * b_vectors[0].toarray().reshape(-1) + \
                            alpha_reg * (A_2_sparse_sqrt.T) @ (A_2_sparse_sqrt @ vertices_concat.numpy()) + \
                            alpha_sharp * A_3_sparse @ vertices_concat.numpy()
        except:
            lin_grad = alpha_0 * A_0_sparse @ vertices_concat.numpy() - alpha_0 * b_vectors[0].toarray().reshape(-1) + alpha_reg * A_2_sparse @ vertices_concat.numpy()
            lin_grad_safe = alpha_0 * (A_0_sparse_sqrt.T) @ (A_0_sparse_sqrt @ vertices_concat.numpy()) + \
                            alpha_0 * (D_e_sparse.T) @ (D_e_sparse @ vertices_concat.numpy()) - \
                            alpha_0 * b_vectors[0].toarray().reshape(-1) + \
                            alpha_reg * (A_2_sparse_sqrt.T) @ (A_2_sparse_sqrt @ vertices_concat.numpy())
        lin_grad = torch.Tensor(lin_grad)
        lin_grad_safe = torch.Tensor(lin_grad_safe)

        lin_grad = torch.cat([lin_grad[:self.n_vertices],
                              lin_grad[self.n_vertices:2 * self.n_vertices],
                              lin_grad[2 * self.n_vertices:]]).to(self.device[0])
        lin_grad_safe = torch.cat([lin_grad_safe[:self.n_vertices],
                                   lin_grad_safe[self.n_vertices:2 * self.n_vertices],
                                   lin_grad_safe[2 * self.n_vertices:]]).to(self.device[0])

        vertices_concat = torch.cat([vertices[:self.n_vertices, 0],
                                     vertices[:self.n_vertices, 1],
                                     vertices[:self.n_vertices, 2]]).double().clone().to(self.device[0]).requires_grad_(True)
        vertices_concat_target = torch.cat([vertices_target[:self.n_vertices, 0],
                                            vertices_target[:self.n_vertices, 1],
                                            vertices_target[:self.n_vertices, 2]]).double().clone()
        vertices_concat_init = vertices_concat.cpu().detach().clone()

        if use_precond:
            time1 = time.time()
            if self.precond is None:
                if load_hessian_path is None:
                    msize = A_matrices[0].shape[0]
                    # A_1 = A_matrices[0][:msize//3, :msize//3].tocsc()
                    # A_1_inv = scipy.sparse.linalg.inv(A_1)

                    A_1 = (self.n_vertices * A_0_sparse)[:msize//3, :msize//3].tocsc()
                    A_1_inv = scipy.sparse.linalg.inv(A_1)
                else:
                    with open(os.path.join(load_hessian_path, 'A_1_inv.pkl'), 'rb') as f:
                        A_1_inv = pickle.load(f)
                A_inv = A_1_inv.tocoo()

                if hessian_cpu:
                    values = A_inv.data
                    indices = np.vstack((A_inv.row, A_inv.col))

                    i = torch.LongTensor(indices)
                    v = torch.DoubleTensor(values)
                    shape = A_inv.shape

                    precond = torch.sparse_coo_tensor(i, v, torch.Size(shape)).double()
                else:
                    A_1_inv = A_1_inv.tocoo()

                    values = A_1_inv.data
                    indices = np.vstack((A_1_inv.row, A_1_inv.col))
                    i = torch.LongTensor(indices)
                    v = torch.DoubleTensor(values)
                    shape = A_1_inv.shape
                    precond_1 = torch.sparse.DoubleTensor(i, v, torch.Size(shape)).to(self.device[0])

                    precond = precond_1

                self.precond = precond
                self.indices = i
                self.values = v
                self.dense_shape = shape

            time2 = time.time()
            optimizer = LBFGS_P([vertices_concat], lr=lr, max_iter=2, history_size=1000, precond=self.precond, hessian_cpu=hessian_cpu,
                                indices=self.indices, values=self.values, dense_shape=self.dense_shape)

            if load_hessian_path is None:
                print('Hessian computing time:', time2 - time1, 's')
            else:
                print('Hessian loading time:', time2 - time1, 's')
        else:
            optimizer = torch.optim.LBFGS([vertices_concat], lr=lr, max_iter=2, history_size=1000)

        history_loss = []
        history_grad_norm = []
        history_data_grad_norm = []
        history_quad_loss = []
        history_smooth_loss = []
        history_p_deviation = []
        history_p_deviation_target = []
        history_p_deviation_mean = []
        history_p_deviation_target_mean = []
        history_data_deviation = []

        data_grad_norm = None

        if self.mapping == 'nn':
            num_parts = 1
        if self.mapping == 'p2p':
            num_parts = len(self.surface_samples_p2p)

        count_close_vertices_track = [[] for _ in range(num_parts)]

        # grad_data = True
        for k in tqdm(range(iterations)):
            # if k % 2 == 0:
            #     grad_data = True
            # else:
            #     grad_data = False

            def closure():

                time1 = time.time()

                nonlocal history_loss
                nonlocal history_data_grad_norm
                nonlocal history_data_deviation
                nonlocal lin_grad
                nonlocal lin_grad_safe
                nonlocal vertices_concat
                nonlocal data_grad_norm
                nonlocal vertices
                nonlocal count_close_vertices_track
                nonlocal num_parts
                # nonlocal grad_data

                optimizer.zero_grad()

                cost_data = 0

                v_coords = torch.cat([vertices_concat[:self.n_vertices][..., None],
                                      vertices_concat[self.n_vertices:2*self.n_vertices][..., None],
                                      vertices_concat[2*self.n_vertices:][..., None]], dim=1).double()

                for i in range(num_parts):
                    if self.mapping == 'nn':
                        cost_data += 1e5 * alpha_data_nn * (torch.sum((self.segmentation_soft_indicators_nn(v_coords[self.surface_samples_nn])) ** 2)) / len(self.surface_samples_nn)
                    elif self.mapping == 'p2p':
                        if self.segmentation_soft_indicators_p2p[i] != 0:
                            len_samples = sum([len(x) for x in self.surface_samples_p2p])
                            kernel_value, count_close_vertices = self.segmentation_soft_indicators_p2p[i](v_coords[self.surface_samples_p2p[i]])
                            count_close_vertices_track[i] += [count_close_vertices]
                            cost_data += alpha_data_p2p * torch.sum(kernel_value ** 2) / len_samples
                    else:
                        raise ValueError('This mapping does not exist')

                time2 = time.time()

                cost_data.backward()

                time3 = time.time()

                data_grad_norm = np.sqrt(torch.sum(vertices_concat.grad.cpu().detach() ** 2).numpy())
                history_data_grad_norm += [data_grad_norm]
                vertices_concat.grad += alpha_quad * lin_grad_safe.double()
                # if not grad_data:
                #     vertices_concat.grad = alpha_quad * lin_grad_safe.double()

                vertices_concat_detached = vertices_concat.cpu().detach()
                try:
                    lin_grad = alpha_0 * A_0_sparse @ vertices_concat_detached.numpy() - alpha_0 * b_vectors[0].toarray().reshape(-1) + alpha_reg * A_2_sparse @ vertices_concat_detached.numpy() + alpha_sharp * A_3_sparse @ vertices_concat_detached.numpy()
                    lin_grad_safe = alpha_0 * (A_0_sparse_sqrt.T) @ (A_0_sparse_sqrt @ vertices_concat_detached.numpy()) + \
                                    alpha_0 * (D_e_sparse.T) @ (D_e_sparse @ vertices_concat_detached.numpy()) - \
                                    alpha_0 * b_vectors[0].toarray().reshape(-1) + \
                                    alpha_reg * (A_2_sparse_sqrt.T) @ (A_2_sparse_sqrt @ vertices_concat_detached.numpy()) + \
                                    alpha_sharp * A_3_sparse @ vertices_concat_detached.numpy()
                except:
                    lin_grad = alpha_0 * A_0_sparse @ vertices_concat_detached.numpy() - alpha_0 * b_vectors[0].toarray().reshape(-1) + alpha_reg * A_2_sparse @ vertices_concat_detached.numpy()
                    lin_grad_safe = alpha_0 * (A_0_sparse_sqrt.T) @ (A_0_sparse_sqrt @ vertices_concat_detached.numpy()) + \
                                    alpha_0 * (D_e_sparse.T) @ (D_e_sparse @ vertices_concat_detached.numpy()) - \
                                    alpha_0 * b_vectors[0].toarray().reshape(-1) + \
                                    alpha_reg * (A_2_sparse_sqrt.T) @ (A_2_sparse_sqrt @ vertices_concat_detached.numpy())
                lin_grad = torch.Tensor(lin_grad)
                lin_grad_safe = torch.Tensor(lin_grad_safe)

                lin_grad = torch.cat([lin_grad[:self.n_vertices],
                                      lin_grad[self.n_vertices:2 * self.n_vertices],
                                      lin_grad[2 * self.n_vertices:]]).to(self.device[0])
                lin_grad_safe = torch.cat([lin_grad_safe[:self.n_vertices],
                                           lin_grad_safe[self.n_vertices:2 * self.n_vertices],
                                           lin_grad_safe[2 * self.n_vertices:]]).to(self.device[0])

                history_loss += [float(cost_data.cpu().detach().numpy())]
                if self.mapping != 'p2p':
                    history_data_deviation += [np.sqrt(np.sum((v_coords[self.surface_samples_nn].cpu().detach().numpy()-self.voxel_centers_nn) ** 2))]

                time4 = time.time()

                # print('Cost computation', time2 - time1)
                # print('Backward', time3 - time2)
                # print('Postcomputation', time4 - time3)
                # print()

                return cost_data

            time5 = time.time()
            if use_precond:
                optimizer.step(closure, self.cuda, self.device_mode)
            else:
                optimizer.step(closure)
            time6 = time.time()
            # print('Step', time6 - time5)

            vertices_concat_detached = vertices_concat.cpu().detach()

            energy = (vertices_concat_detached.numpy().T @ A_0_sparse @ vertices_concat_detached.numpy() - 2 * b_vectors[0].toarray().reshape(-1).T @ vertices_concat_detached.numpy() + torch.sum(affine_components ** 2).numpy() + torch.sum(translation_components ** 2).numpy()) / self.n_vertices
            energy_0 = (vertices_concat_detached.numpy().T @ A_0_sparse @ vertices_concat_detached.numpy() - 2 * b_vectors[0].toarray().reshape(-1).T @ vertices_concat_detached.numpy()) / self.n_vertices
            energy_0_safe = ((A_0_sparse_sqrt @ vertices_concat_detached.numpy()).T @ (A_0_sparse_sqrt @ vertices_concat_detached.numpy()) +
                             (D_e_sparse @ vertices_concat_detached.numpy()).T @ (D_e_sparse @ vertices_concat_detached.numpy()) -
                             2 * b_vectors[0].toarray().reshape(-1).T @ vertices_concat_detached.numpy()) / self.n_vertices

            energy_smooth = (vertices_concat_detached.numpy().T @ A_2_sparse @ vertices_concat_detached.numpy()) / self.n_vertices
            energy_smooth_safe = ((A_2_sparse_sqrt @ vertices_concat_detached.numpy()).T @ (A_2_sparse_sqrt @ vertices_concat_detached.numpy())) / self.n_vertices

            grad_norm = np.sqrt(torch.sum(lin_grad.cpu() ** 2).numpy())
            grad_norm_safe = np.sqrt(torch.sum(lin_grad_safe.cpu() ** 2).numpy())
            history_grad_norm += [grad_norm]
            history_quad_loss += [energy]
            history_smooth_loss += [energy_smooth]
            p_deviation = torch.sqrt(torch.sum((vertices_concat_init - vertices_concat_detached) ** 2))
            history_p_deviation += [p_deviation]
            p_deviation_target = torch.sqrt(torch.sum((vertices_concat_target - vertices_concat_detached) ** 2))
            history_p_deviation_target += [p_deviation_target]

            v_coords_init = vertices_concat_init.numpy().reshape((-1, 3), order='F')
            v_coords_detached = vertices_concat_detached.numpy().reshape((-1, 3), order='F')
            v_coords_target = vertices_concat_target.numpy().reshape((-1, 3), order='F')
            p_deviation_mean = np.sqrt(np.sum((v_coords_init - v_coords_detached) ** 2, axis=1).mean())
            history_p_deviation_mean += [p_deviation_mean]
            p_deviation_target_mean = np.sqrt(np.sum((v_coords_target - v_coords_detached) ** 2, axis=1).mean())
            history_p_deviation_target_mean += [p_deviation_target_mean]

            if k % print_freq == 0:
                print(
                    'p_dev {p_dev:3.3f} '
                    'E0 {E0:3.3f} E0_safe {E0_safe:3.3f} ES {ES:3.3f} ES_safe {ES_safe:3.3f} ED {ED:3.3f} '
                    'LG_norm {LG_norm:3.3f} LG_norm_safe {LG_norm_safe:3.3f} '
                    'DG_norm {DG_norm:3.3f}'.format(
                    **{
                        'timestamp': datetime.now().strftime("%d.%m.%Y %H:%M:%S.%f"),
                        'p_dev': p_deviation.numpy(),
                        'E0': energy_0,
                        'E0_safe': energy_0_safe,
                        'ES': energy_smooth,
                        'ES_safe': energy_smooth_safe,
                        'ED': history_loss[-1],
                        'LG_norm': grad_norm,
                        'LG_norm_safe': grad_norm_safe,
                        'DG_norm': data_grad_norm
                    }
                ), file=sys.stdout)
                sys.stdout.flush()

            time7 = time.time()
            # print('Printing', time7 - time6)
            # print()

        vertices_concat = vertices_concat.cpu().detach().numpy().reshape((-1, 3), order='F')
        self.vertices_approximation = vertices_concat

        if plot:
            plot_info(history_grad_norm, history_quad_loss,
                      history_smooth_loss, history_loss,
                      history_p_deviation, history_p_deviation_target,
                      history_p_deviation_mean, history_p_deviation_target_mean)

        if self.mapping != 'p2p':
            return history_data_deviation
        else:
            return history_data_deviation, count_close_vertices_track, vertices_concat_init, vertices_concat_detached

    def get_new_parts_vertices(self):

        parts_vertices_pred_numpy = [x.detach().numpy()[:, :3] for x in self.parts_vertices_pred]
        return parts_vertices_pred_numpy

    def get_starting_parts_vertices(self):

        parts_vertices_pred_numpy = [x.numpy()[:, :3] for x in self.parts_starting_vertices]
        return parts_vertices_pred_numpy

    def get_initial_parts_vertices(self):

        parts_vertices_initial_numpy = self.target_vertices.numpy()[:, :3]
        return parts_vertices_initial_numpy

    def get_indicator_functions(self):

        return self.segmentation_soft_indicators

    def get_bitriangles_map(self):

        return self.bitriangles_map

    def get_first_approximation(self):

        return self.vertices_approximation

    def print_times(self):

        print('Deformations time:', self.deformations_time, 's')
        print('Cost time:', self.cost_time, 's')
        print('Backward time:', self.backward_time, 's')

    def get_target_deformations(self):

        return self.target_edges_deformations


import pymongo
import json
import argparse
import traceback
import datetime
import pytz
from astropy.time import Time
import tqdm
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import normalize
from bson.json_util import loads, dumps
import numpy as np
import pandas as pd
import gzip
import io
from astropy.io import fits

''' load config and secrets '''
with open('/app/config.json') as cjson:
    config = json.load(cjson)

with open('/app/secrets.json') as sjson:
    secrets = json.load(sjson)

for k_ in secrets:
    config[k_].update(secrets.get(k_, {}))

def utc_now():
    return datetime.datetime.now(pytz.utc)

def time_stamps():
    """

    :return: local time, UTC time
    """
    return datetime.datetime.now().strftime('%Y%m%d_%H:%M:%S'), \
           datetime.datetime.utcnow().strftime('%Y%m%d_%H:%M:%S')

def connect_to_db():
    """ Connect to the mongodb database

    :return:
    """
    try:
        # there's only one instance of DB, it's too big to be replicated
        _client = pymongo.MongoClient(host=config['database']['host'],
                                      port=config['database']['port'])
        # grab main database:
        _db = _client[config['database']['db']]
    except Exception as _e:
        raise ConnectionRefusedError
    try:
        # authenticate
        _db.authenticate(config['database']['user'], config['database']['pwd'])
    except Exception as _e:
        raise ConnectionRefusedError

    return _client, _db

def insert_db_entry(_db, _collection=None, _db_entry=None):
    """
        Insert a document _doc to collection _collection in DB.
        It is monitored for timeout in case DB connection hangs for some reason
    :param _collection:
    :param _db_entry:
    :return:
    """
    assert _collection is not None, 'Must specify collection'
    assert _db_entry is not None, 'Must specify document'
    try:
        _db[_collection].insert_one(_db_entry)
    except Exception as _e:
        print('Error inserting {:s} into {:s}'.format(str(_db_entry['_id']), _collection))
        traceback.print_exc()
        print(_e)

def insert_multiple_db_entries(_db, _collection=None, _db_entries=None, _verbose=False):
    """
        Insert a document _doc to collection _collection in DB.
        It is monitored for timeout in case DB connection hangs for some reason
    :param _db:
    :param _collection:
    :param _db_entries:
    :param _verbose:
    :return:
    """
    assert _collection is not None, 'Must specify collection'
    assert _db_entries is not None, 'Must specify documents'
    try:
        _db[_collection].insert_many(_db_entries, ordered=False)
    except pymongo.errors.BulkWriteError as bwe:
        if _verbose:
            print(bwe.details)
    except Exception as _e:
        if _verbose:
            traceback.print_exc()
            print(_e)

def make_triplet(alert, to_tpu: bool = False):
    """
        Feed in alert packet
    """
    cutout_dict = dict()

    for cutout in ('science', 'template', 'difference'):
        cutout_data = loads(dumps([alert[f'cutout{cutout.capitalize()}']['stampData']]))[0]

        # unzip
        with gzip.open(io.BytesIO(cutout_data), 'rb') as f:
            with fits.open(io.BytesIO(f.read())) as hdu:
                data = hdu[0].data
                # replace nans with zeros
                cutout_dict[cutout] = np.nan_to_num(data)
                # normalize
                # cutout_dict[cutout] = normalize(cutout_dict[cutout])
                # L2-normalize
                cutout_dict[cutout] /= np.linalg.norm(cutout_dict[cutout])

        # pad to 63x63 if smaller
        shape = cutout_dict[cutout].shape
        if shape != (63, 63):
            # print(f'Shape of {candid}/{cutout}: {shape}, padding to (63, 63)')
            cutout_dict[cutout] = np.pad(cutout_dict[cutout], [(0, 63 - shape[0]), (0, 63 - shape[1])],
                                         mode='constant', constant_values=1e-9)

    triplet = np.zeros((63, 63, 3))
    triplet[:, :, 0] = cutout_dict['science']
    triplet[:, :, 1] = cutout_dict['template']
    triplet[:, :, 2] = cutout_dict['difference']

    if to_tpu:
        # Edge TPUs require additional processing
        triplet = np.rint(triplet * 128 + 128).astype(np.uint8).flatten()

    return triplet

def alert_filter__ml(alert, ml_models: dict = None):
    """Filter to apply to each alert.
    """

    scores = dict()

    try:
        ''' braai '''
        triplet = make_triplet(alert)
        triplets = np.expand_dims(triplet, axis=0)
        braai = ml_models['braai']['model'].predict(x=triplets)[0]
        # braai = 1.0
        scores['braai'] = float(braai)
        scores['braai_version'] = ml_models['braai']['version']
    except Exception as e:
        print(*time_stamps(), str(e))

    return scores

def main(obs_date=datetime.datetime.utcnow().strftime('%Y%m%d')):

    # ML models:
    ml_models = dict()
    for m in config['ml_models']:
        try:
            m_v = config["ml_models"][m]["version"]
            ml_models[m] = {'model': load_model(f'/app/models/{m}_{m_v}.h5'),
                            'version': m_v}
        except Exception as e:
            print(*time_stamps(), f'Error loading ML model {m}')
            traceback.print_exc()
            print(e)
            continue

    jd = Time(datetime.datetime.strptime(obs_date, '%Y%m%d')).jd

    collection_alerts = 'ZTF_alerts'

    # connect to MongoDB:
    print('Connecting to DB')
    client, db = connect_to_db()
    print('Successfully connected')

    # db[collection_alerts].create_index([('candidate.jd', pymongo.DESCENDING),
    #                                     ('classifications.braai', pymongo.DESCENDING),
    #                                     ('candidate.rb', pymongo.DESCENDING),
    #                                     ('candid', pymongo.DESCENDING)],
    #                                    background=True)

    num_doc = db[collection_alerts].count_documents({'candidate.jd': {'$gt': jd, '$lt': jd+1}})
    print(num_doc)

    cursor = db[collection_alerts].find({'candidate.jd': {'$gt': jd, '$lt': jd+1}},
                                        {'candidate': 0, 'prv_candidates': 0, 'coordinates': 0})

    # for alert in cursor.limit(1):
    for alert in tqdm.tqdm(cursor, total=num_doc):
        scores = alert_filter__ml(alert, ml_models)
        # print(alert['candid'], scores)
        db[collection_alerts].update_one({'_id': alert['_id']},
                                         {'$set': {'classifications': scores}})

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Apply ML models on a night worth of ZTF alerts')
    parser.add_argument('--obsdate', help='observing date', default=datetime.datetime.utcnow().strftime('%Y%m%d'))

    args = parser.parse_args()

    main(obs_date=args.obsdate)


import operator
import typing
from abc import ABC
from collections import Counter
from enum import Enum
from itertools import chain, product
from typing import Callable

import numpy as np
import SimpleITK as sitk
from napari.layers import Image
from napari.types import LayerDataTuple

from PartSegCore.algorithm_describe_base import AlgorithmDescribeBase, AlgorithmProperty, ROIExtractionProfile
from PartSegCore.analysis.measurement_calculation import Diameter, get_border
from PartSegCore.autofit import density_mass_center
from PartSegCore.channel_class import Channel
from PartSegCore.class_generator import enum_register
from PartSegCore.roi_info import ROIInfo
from PartSegCore.segmentation import RestartableAlgorithm
from PartSegCore.segmentation.algorithm_base import AdditionalLayerDescription, SegmentationResult
from PartSegCore.segmentation.noise_filtering import noise_filtering_dict
from PartSegCore.segmentation.threshold import BaseThreshold, threshold_dict

from .widgets import TrapezoidWidget

class NeuType(Enum):
    PMN_neu = 1
    RND_neu = 2
    NER_neu = 3
    PMP_neu = 4
    Bacteria = 5
    NET = 8
    Unknown_intra = 10
    Unknown_extra = 11

    def __str__(self):
        return self.name.replace("_", " ")

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

    @classmethod
    def known_components(cls):
        return (x for x in cls.__members__.values() if x.value < 10)

    @classmethod
    def neutrofile_components(cls):
        return (x for x in cls.__members__.values() if x.value < 8)

    @classmethod
    def all_components(cls):
        return cls.__members__.values()

try:
    # noinspection PyUnresolvedReferences,PyUnboundLocalVariable
    reloading
except NameError:
    reloading = False  # means the module is being imported
    enum_register.register_class(NeuType)

LABELING_NAME = "Labeling"
SCORE_SUFFIX = "_score"
CATEGORY_STR = "category"
PARAMETER_TYPE_LIST = [
    "pixel count",
    "brightness",
    "ext. brightness",
    "brightness gradient",
]  # "brightness", "roundness"
DESCRIPTION_DICT = {
    NeuType.PMN_neu: "polymorphonuclear neutrophils",
    NeuType.RND_neu: "decondensed chromatin neutrophils",
    NeuType.NER_neu: "ruptured nuclear envelope neutrophils",
    NeuType.PMP_neu: "plasma membrane permeabilized neutrophils",
    NeuType.NET: "neutrophil extracellular trap",
    NeuType.Unknown_intra: "unclassified intracellurar component",
    NeuType.Unknown_extra: "unclassified extracellurar component",
    NeuType.Bacteria: "bacteria",
}

class NeutrofileSegmentationBase(RestartableAlgorithm, ABC):
    @classmethod
    def support_time(cls):
        return False

    @classmethod
    def support_z(cls):
        return False

    def _calculate_mask(self, channel, threshold_name):
        thr: BaseThreshold = threshold_dict[self.new_parameters[threshold_name]["name"]]
        mask, thr_val = thr.calculate_mask(
            channel, self.mask, self.new_parameters[threshold_name]["values"], operator.gt
        )
        return mask, thr_val

    @staticmethod
    def _calc_components(arr, min_size, close_holes=0):
        components = sitk.RelabelComponent(
            sitk.ConnectedComponent(sitk.GetImageFromArray(arr)),
            min_size,
        )

        if close_holes:
            binary1 = sitk.BinaryThreshold(components, 0, 0)
            removed_small = sitk.RelabelComponent(sitk.ConnectedComponent(binary1), close_holes)
            binary2 = sitk.BinaryThreshold(removed_small, 0, 0)
            components = sitk.RelabelComponent(
                sitk.ConnectedComponent(binary2),
                min_size,
            )
        return sitk.GetArrayFromImage(components)

class Trapalyzer(NeutrofileSegmentationBase):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.count_dict = Counter()
        self.area_dict = Counter()
        self.nets = 0
        self.other = 0
        self.net_size = 0
        self.quality = 0
        self.image_size = 0

    def classify_nets(self, outer_dna_mask, cleaned_inner, cleaned_outer, net_size):
        nets = self._calc_components(outer_dna_mask, int(net_size["lower_bound"]))
        laplacian_outer_image = _laplacian_estimate(cleaned_outer, 1.3)
        annotation = {}
        i = 1
        nets[nets > 0] += 1
        for val in np.unique(nets):
            if val == 0:
                continue
            component = np.array(nets == val)
            brightness_gradient = np.mean(laplacian_outer_image[component])
            # brightness_gradient_score = sine_score_function(
            #     brightness_gradient,
            #     softness=self.new_parameters["softness"],
            #     **self.new_parameters["net_brightness_gradient"],
            # )
            brightness_std = np.std(cleaned_outer[component])
            brightness_std_score = sine_score_function(
                np.std(cleaned_outer[component]),
                softness=self.new_parameters["softness"],
                **self.new_parameters["net_ext_brightness_std"],
            )
            brightness = np.quantile(cleaned_inner[component], 0.9)
            ext_brightness = np.quantile(cleaned_outer[component], 0.9)
            ext_brightness_score = sine_score_function(
                ext_brightness, softness=self.new_parameters["softness"], **self.new_parameters["net_ext_brightness"]
            )
            voxels = np.count_nonzero(component)
            voxels_score = sine_score_function(
                voxels, softness=self.new_parameters["softness"], **self.new_parameters["net_size"]
            )

            if voxels_score * ext_brightness_score * brightness_std_score < self.new_parameters["minimum_score"]:
                if not self.new_parameters["unknown_net"]:
                    nets[component] = 0
                    continue
                else:
                    category = NeuType.Unknown_extra
            else:
                category = NeuType.NET

            data_dict = {
                "component_id": i,
                CATEGORY_STR: category,
                "pixel count": voxels,
                "brightness": brightness,
                "ext. brightness": ext_brightness,
                # "brightness gradient": np.mean(laplacian_image[component]),
                "ext. brightness gradient": brightness_gradient,
                "ext. brightness SD": brightness_std,
            }
            annotation[i] = data_dict
            nets[component] = i
            i += 1
        return nets, annotation

    def calculation_run(self, report_fun: Callable[[str, int], None]) -> SegmentationResult:
        self.count_dict = Counter()
        self.area_dict = Counter()
        self.nets = 0
        self.other = 0
        self.net_size = 0
        inner_dna_channel = self.get_channel(self.new_parameters["inner_dna"])
        self.image_size = inner_dna_channel.shape[0] * inner_dna_channel.shape[1] * inner_dna_channel.shape[2]
        inner_noise_filtering_parameters = self.new_parameters["inner_dna_noise_filtering"]
        cleaned_inner = noise_filtering_dict[inner_noise_filtering_parameters["name"]].noise_filter(
            inner_dna_channel, self.image.spacing, inner_noise_filtering_parameters["values"]
        )
        inner_dna_mask, thr_val = self._calculate_mask(cleaned_inner, "inner_threshold")

        outer_dna_channel = self.get_channel(self.new_parameters["outer_dna"])
        outer_noise_filtering_parameters = self.new_parameters["outer_dna_noise_filtering"]
        cleaned_outer = noise_filtering_dict[outer_noise_filtering_parameters["name"]].noise_filter(
            outer_dna_channel, self.image.spacing, outer_noise_filtering_parameters["values"]
        )
        outer_dna_mask, dead_thr_val = self._calculate_mask(cleaned_outer, "outer_threshold")
        outer_dna_components, net_annotation = self.classify_nets(
            outer_dna_mask, cleaned_inner, cleaned_outer, self.new_parameters["net_size"]
        )
        inner_dna_mask[outer_dna_components > 0] = 0
        size_param_array = [
            self.new_parameters[f"{x.name.lower()}_pixel count"] for x in NeuType.neutrofile_components()
        ]

        min_object_size = int(
            max(
                5,
                min(
                    x["lower_bound"]
                    - (x["upper_bound"] - x["lower_bound"])
                    * self.new_parameters["softness"]
                    * x["lower_bound"]
                    / (x["upper_bound"] - x["lower_bound"])
                    for x in size_param_array
                ),
                self.new_parameters["minimum_size"],
            )
        )
        inner_dna_components = self._calc_components(inner_dna_mask, min_object_size, close_holes=100)

        result_labeling, roi_annotation = self._classify_neutrofile(inner_dna_components, cleaned_inner, cleaned_outer)

        nets_components = [k for k, v in net_annotation.items() if v[CATEGORY_STR] == NeuType.NET]
        unknown_net_components = [k for k, v in net_annotation.items() if v[CATEGORY_STR] != NeuType.NET]

        self.count_dict[NeuType.NET] = len(nets_components)

        self.count_dict[NeuType.Unknown_extra] = len(unknown_net_components)

        result_labeling[np.isin(outer_dna_components, nets_components)] = NeuType.NET.value
        result_labeling[np.isin(outer_dna_components, unknown_net_components)] = NeuType.Unknown_extra.value
        max_component = inner_dna_components.max()
        inner_dna_components[outer_dna_components > 0] = outer_dna_components[outer_dna_components > 0] + max_component
        for value in np.unique(inner_dna_components[outer_dna_components > 0]):
            roi_annotation[value] = net_annotation[value - max_component]
            roi_annotation[value]["component_id"] = value
        alternative_representation = {LABELING_NAME: result_labeling}
        for val in NeuType.all_components():
            alternative_representation[str(val)] = (result_labeling == val.value).astype(np.uint8) * val.value
        self.net_size = np.count_nonzero(alternative_representation[str(NeuType.NET)])
        self.area_dict[NeuType.NET] = self.net_size

        from .measurement import QualityMeasure

        self.quality = QualityMeasure.calculate_property(inner_dna_components, roi_annotation)
        return SegmentationResult(
            inner_dna_components,
            self.get_segmentation_profile(),
            alternative_representation=alternative_representation,
            roi_annotation=roi_annotation,
            additional_layers={
                "inner_mask": AdditionalLayerDescription(inner_dna_mask.astype(np.uint8), "labels", "inner_mask"),
                "cleaned outer brightness gradient": AdditionalLayerDescription(
                    _laplacian_estimate(cleaned_outer, 1.3), "image", "cleaned outer brightness gradient"
                ),
                "cleaned inner brightness gradient": AdditionalLayerDescription(
                    _laplacian_estimate(cleaned_inner, 1.3), "image", "cleaned inner brightness gradient"
                ),
            },
        )

    def _classify_neutrofile(self, inner_dna_components, cleaned_inner, cleaned_outer):
        self.cell_count = 0
        laplacian_image = _laplacian_estimate(cleaned_inner, 1.3)
        annotation = {}
        labeling = np.zeros(inner_dna_components.shape, dtype=np.uint16)
        bounds = ROIInfo.calc_bounds(inner_dna_components)
        inner_dna_components_border = get_border(inner_dna_components)

        for val in np.unique(inner_dna_components):
            if val == 0:
                continue
            slices = tuple(bounds[val].get_slices(margin=5))
            component = np.array(inner_dna_components[slices] == val)
            component_border = inner_dna_components_border[slices] == val
            voxels = np.count_nonzero(component)
            perimeter = np.count_nonzero(component_border)
            diameter = Diameter.calculate_property(component, voxel_size=(1,) * component.ndim, result_scalar=1)

            if perimeter == 0:
                continue
            data_dict = {
                "pixel count": voxels,
                "brightness gradient": np.mean(laplacian_image[slices][component]),
                "brightness": np.quantile(cleaned_inner[slices][component], 0.9),
                "intensity": np.sum(cleaned_inner[slices][component]),
                # "homogenity": np.mean(inner_dna_channel[component]) / np.std(inner_dna_channel[component]),
                "ext. brightness": np.quantile(cleaned_outer[slices][component], 0.9),
                #                  "roundness": new_sphericity(component, self.image.voxel_size),
                #                 "roundness2": voxels / ((diameter ** 2 / 4) * pi),
                #                 "diameter": diameter,
                # "curvature": np.std(curvature(component_border_coords[0], component_border_coords[1])),
                #                 "border_size": component_border_coords[0].size,
                # "circumference": len(component_border_coords[0]),
                # "area to circumference": voxels / len(component_border_coords[0]),
                # "signal colocalization": colocalization1,
                "circularity": 3 * np.pi * voxels / perimeter**2,  # inspired by NETQUANT approach
                # "circularity2": voxels / ((diameter ** 2 / 4) * np.pi),
            }
            annotation[val] = dict(
                {"component_id": val, CATEGORY_STR: NeuType.Unknown_intra},
                **data_dict,
                **{
                    f"{prefix} {suffix}": sine_score_function(
                        data_dict[suffix],
                        softness=self.new_parameters["softness"],
                        **self.new_parameters[f"{prefix.name.lower()}_{suffix}"],
                    )
                    for prefix, suffix in product(NeuType.neutrofile_components(), PARAMETER_TYPE_LIST)
                },
            )
            score_list = []
            for component_name in NeuType.neutrofile_components():
                score = 1
                for parameter in PARAMETER_TYPE_LIST:
                    score *= annotation[val][f"{component_name} {parameter}"]
                score_list.append((score, component_name))
                annotation[val][str(component_name) + SCORE_SUFFIX] = score

            score_list = sorted(score_list)
            if (
                score_list[-1][0] < self.new_parameters["minimum_score"]
                or score_list[-2][0] > self.new_parameters["maximum_other"]
            ):
                labeling[inner_dna_components == val] = NeuType.Unknown_intra.value
                self.count_dict[NeuType.Unknown_intra] += 1
            else:
                labeling[inner_dna_components == val] = score_list[-1][1].value
                self.count_dict[score_list[-1][1]] += 1
                self.cell_count += 1
                self.area_dict[score_list[-1][1]] += voxels
                annotation[val][CATEGORY_STR] = score_list[-1][1]

        return labeling, annotation

    def get_info_text(self):
        image_size = self.image_size if self.image_size else 1
        return (
            f"Annotation quality: {self.quality}"
            + "\n"
            + "\n".join(
                f"{name}: {self.count_dict[name]}, {self.area_dict[name]} px"
                f" ({round(100*self.area_dict[name]/image_size, 2)}% image area)"
                for name in [NeuType.NET, NeuType.Unknown_extra]
            )
            + "\n"
            + "\n".join(
                f"{name}: {self.count_dict[name]} ({round(100*self.count_dict[name] / (self.cell_count or 1), 2)}"
                f" % cells), {self.area_dict[name]} px ({round(100*self.area_dict[name]/image_size, 2)}%"
                f" image area)"
                for name in NeuType.neutrofile_components()
            )
        )

    def get_segmentation_profile(self) -> ROIExtractionProfile:
        return ROIExtractionProfile("", self.get_name(), dict(self.new_parameters))

    @classmethod
    def get_name(cls) -> str:
        return "Trapalyzer"

    @classmethod
    def get_fields(cls) -> typing.List[typing.Union[AlgorithmProperty, str]]:
        initial = list(
            chain(
                *(
                    [
                        AlgorithmProperty(
                            f"{prefix.name.lower()}_{suffix}",
                            f"{prefix} {suffix}",
                            {"lower_bound": 0, "upper_bound": 1},
                            property_type=TrapezoidWidget,
                            help_text=f"{DESCRIPTION_DICT[prefix]} {suffix}",
                        )
                        for suffix in PARAMETER_TYPE_LIST
                    ]
                    + ["-------------------"]
                    for prefix in NeuType.neutrofile_components()
                )
            )
        ) + [
            AlgorithmProperty("minimum_score", "Minimum score", 0.8),
            AlgorithmProperty("maximum_other", "Maximum competing score", 0.4),
            AlgorithmProperty(
                "minimum_size", "Minimum comp. pixel count", 40, (1, 9999), help_text="Minimum component pixel count"
            ),
            AlgorithmProperty("softness", "Error margin coeficient", 0.1, (0, 1)),
        ]

        thresholds = [
            AlgorithmProperty("inner_dna", "All DNA channel", 1, property_type=Channel),
            AlgorithmProperty(
                "inner_dna_noise_filtering",
                "Filter type",
                next(iter(noise_filtering_dict.keys())),
                possible_values=noise_filtering_dict,
                value_type=AlgorithmDescribeBase,
            ),
            AlgorithmProperty(
                "inner_threshold",
                "Threshold type",
                next(iter(threshold_dict.keys())),
                possible_values=threshold_dict,
                property_type=AlgorithmDescribeBase,
            ),
            AlgorithmProperty("outer_dna", "Extracellural DNA channel", 1, property_type=Channel),
            AlgorithmProperty(
                "outer_dna_noise_filtering",
                "Filter type",
                next(iter(noise_filtering_dict.keys())),
                possible_values=noise_filtering_dict,
                value_type=AlgorithmDescribeBase,
            ),
            AlgorithmProperty(
                "outer_threshold",
                "Threshold type",
                next(iter(threshold_dict.keys())),
                possible_values=threshold_dict,
                property_type=AlgorithmDescribeBase,
            ),
            AlgorithmProperty(
                "net_size",
                "NET pixel count",
                {"lower_bound": 850, "upper_bound": 999999},
                property_type=TrapezoidWidget,
            ),
            AlgorithmProperty(
                "net_ext_brightness",
                "NET extracellular brightness",
                {"lower_bound": 21, "upper_bound": 100},
                property_type=TrapezoidWidget,
            ),
            # AlgorithmProperty(
            #     "net_brightness_gradient",
            #     "NET ext. brightness gradient",
            #     {"lower_bound": -1.0, "upper_bound": 1.0},
            #     property_type=TrapezoidWidget,
            # ),
            AlgorithmProperty(
                "net_ext_brightness_std",
                "NET ext. brightness SD",
                {"lower_bound": 0.0, "upper_bound": 1.0},
                property_type=TrapezoidWidget,
            ),
            AlgorithmProperty(
                "unknown_net", "detect extracellular artifacts", True, help_text="If mark unknown net components"
            ),
            "-----------------------",
        ]

        return thresholds + initial

def trapezoid_score_function(x, lower_bound, upper_bound, softness=0.5):
    """
    Compute a score on a scale from 0 to 1 that indicate whether values from x belong
    to the interval (lbound, ubound) with a softened boundary.
    If a point lies inside the interval, its score is equal to 1.
    If the point is further away than the interval length multiplied by the softness parameter,
    its score is equal to zero.
    Otherwise the score is given by a linear function.
    """
    interval_width = upper_bound - lower_bound
    subound = upper_bound + softness * interval_width
    slbound = lower_bound - (softness * interval_width) * 0.1
    swidth = softness * interval_width  # width of the soft boundary
    if lower_bound <= x <= upper_bound:
        return 1.0
    elif x <= slbound or x >= subound:
        return 0.0
    elif slbound <= x <= lower_bound:
        return 1.0 - (lower_bound - x) / swidth
    elif upper_bound <= x <= subound:
        return 1.0 - (x - upper_bound) / swidth

def gaussian_score_function(x, lbound, ubound, softness=0.5):
    """
    Computes a score on a scale from 0 to 1 whether x lies within the [lbound, ubound] interval.
    Fuzzified by a gaussian function.
    Softness controls the extension of the interval.
    """
    assert ubound >= lbound, "ubound needs to be >= than lbound"
    interval_mean = ubound + lbound
    interval_length = ubound - lbound
    if lbound <= x <= ubound:
        return 1.0
    sd_l = np.log(softness * interval_length * lbound / interval_mean) - np.log(3)  # to avoid raising to small powers
    sd_l = np.exp(sd_l)
    sd_u = np.log(softness * interval_length * ubound / interval_mean) - np.log(3)
    sd_u = np.exp(sd_u)
    if x <= lbound - 3.5 * sd_l or x >= ubound + 3.5 * sd_u:
        return 0.0
    if lbound - 3.5 * sd_l <= x <= lbound:
        logscore = -0.5 * (lbound - x) ** 2 / sd_l**2
        return np.exp(logscore)
    if ubound <= x <= ubound + 3.5 * sd_u:
        logscore = -0.5 * (x - ubound) ** 2 / sd_u**2
        return np.exp(logscore)

def sine_score_function(x, lower_bound, upper_bound, softness=0.5):
    """
    Computes a score on a scale from 0 to 1 whether x lies within the [lbound, ubound] interval.
    Extended with a sine function.
    Softness controls the extension of the interval.
    """
    assert upper_bound >= lower_bound, "upper_bound needs to be >= than lower_bound"
    interval_mean = upper_bound + lower_bound
    interval_length = upper_bound - lower_bound
    if lower_bound <= x <= upper_bound:
        return 1.0
    extension_span_left = softness * interval_length * lower_bound / interval_mean
    extension_span_right = softness * interval_length * upper_bound / interval_mean
    if x <= lower_bound - extension_span_left or x >= upper_bound + extension_span_right:
        return 0.0
    if lower_bound - extension_span_left <= x <= lower_bound:
        coord_transform = 2 * (x - lower_bound + extension_span_left) / extension_span_left - 1
    elif upper_bound <= x <= upper_bound + extension_span_right:
        coord_transform = 2 * (x - upper_bound + extension_span_right) / extension_span_right - 1
    else:
        raise RuntimeError("Something went terribly wrong!")
    return 0.5 + 0.5 * np.sin(0.5 * np.pi * coord_transform)

def new_sphericity(component: np.ndarray, voxel_size):
    component = component.squeeze()
    voxel_area = np.prod(voxel_size)
    center = np.array([density_mass_center(component, voxel_size)])
    area = np.count_nonzero(component) * voxel_area
    coordinates = np.transpose(np.nonzero(component)) * voxel_size - center
    radius_square = area / np.pi
    return np.sum(np.sum(coordinates**2, axis=1) <= radius_square) / (
        area / voxel_area + np.sum(np.sum(coordinates**2, axis=1) > radius_square)
    )

def laplacian_estimate(image: Image, radius=1.30, clip_bellow_0=True) -> LayerDataTuple:
    res = _laplacian_estimate(image.data[0], radius=radius)
    if clip_bellow_0:
        res[res < 0] = 0
    res = res.reshape(image.data.shape)
    return LayerDataTuple((res, {"colormap": "magma", "scale": image.scale, "name": "Laplacian estimate"}))

def _laplacian_estimate(channel: np.ndarray, radius=1.30) -> np.ndarray:
    data = channel.astype(np.float64).squeeze()
    return -sitk.GetArrayFromImage(sitk.LaplacianRecursiveGaussian(sitk.GetImageFromArray(data), radius)).reshape(
        channel.shape
    )

def curvature(x, y, *args):
    assert len(x) == len(y)
    n = len(x)
    k = np.zeros(n)
    for i in range(n):
        ddx = x[(i + 1) % n] + x[(i - 1) % n] - 2 * x[i]
        ddx *= n**2
        ddy = y[(i + 1) % n] + y[(i - 1) % n] - 2 * y[i]
        ddy *= n**2
        k[i] = np.sqrt(ddx**2 + ddy**2)
    return k


import sys

from collections import Counter
from io import StringIO
import numpy as np
import gym
import gym.spaces

class TimeTable (gym.Env):
  """
  -1 配置不可の枠
  0  未配置の枠
  1  現代文 週3
  2  古典 週2
  3  漢文 週1
  4  数学 週6 教師A
  5  英語 週6 教師A
  6  日本史 週2
  7  世界史 週2
  8  化学 週2
  9  物理 週2
  10 生物 週2
  11 数学 週6 教師B
  12 英語 週6 教師B
  13 体育 週2
  14 美術 週1
  15 書道 週1
  16 保健 週1
  """

  CLASS_LIST = [0, 1, 2, 3, 4, 5]

  CLASS_SIZE = len(CLASS_LIST)

  CLASS_LESSON = {
    1 : [0, 1, 2, 3, 4, 5],
    2 : [0, 1, 2, 3, 4, 5],
    3 : [0, 1, 2, 3, 4, 5],
    4 : [0, 1, 2],
    5 : [0, 1, 2],
    6 : [0, 1, 2],
    7 : [3, 4, 5],
    8 : [0, 1, 2, 3],
    9 : [2, 3, 4, 5],
    10: [0, 3, 4, 5],
    11: [3, 4, 5],
    12: [3, 4, 5],
    13 : [0, 1, 2, 3, 4, 5],
    14 : [0, 1, 2, 3, 4, 5],
    15 : [0, 1, 2, 3, 4, 5],
    16 : [0, 1, 2, 3, 4, 5]
  }

  LESSON_NAME = {
    -1: '禁 ', # 1  現代文 週3
    0 : '未 ', # 1  現代文 週3
    1 : '現0', # 1  現代文 週3
    2 : '古0', # 2  古典 週2
    3 : '漢0', # 3  漢文 週1
    4 : '数0', # 4  数学 週6 教師A
    5 : '英0', # 5  英語 週6 教師A
    6 : '日0', # 6  日本史 週2
    7 : '世0', # 7  世界史 週2
    8 : '化0', # 8  化学 週2
    9 : '物0', # 9  物理 週2
    10: '生0', # 10 生物 週2
    11: '数1', # 11 数学 週6 教師B
    12: '英1', # 12 英語 週6 教師B
    13: '体0', # 13 体育 週2
    14: '美0', # 14 美術 週1
    15: '書0', # 15 書道 週1
    16: '保0'  # 16 保健 週1
  }

  CLASS_LESSON_TIMES = {
    1 : 3,
    2 : 2,
    3 : 1,
    4 : 6,
    5 : 6,
    6 : 2,
    7 : 2,
    8 : 2,
    9 : 2,
    10: 2,
    11: 6,
    12: 6,
    13 : 2,
    14 : 1,
    15 : 1,
    16 : 1
  }

  WEEK_SIZE = 6

  DAY_SIZE = 5

  TIME_TABLE = np.zeros((CLASS_SIZE, WEEK_SIZE * DAY_SIZE))

  LESSON_LIST = CLASS_LESSON.keys()

  LESSON_SIZE = len(LESSON_LIST)

  TOTAL_LESSON = CLASS_SIZE * WEEK_SIZE * DAY_SIZE

  MAX_DAMAGE = 100

  def __init__(self):
    super().__init__()
    # 土曜日の授業のコマ数
    for i in range(self.CLASS_SIZE):
      self.TIME_TABLE[i][-1] = -1
      self.TIME_TABLE[i][-2] = -1
    self.action_space = gym.spaces.Discrete(self.CLASS_SIZE * self.WEEK_SIZE * self.DAY_SIZE * self.LESSON_SIZE)
    self.observation_space = gym.spaces.Box(
        low = -1,
        high = max(self.LESSON_LIST),
        shape = self.TIME_TABLE.shape
    )
    # self.reward_range = [0, self.TOTAL_LESSON]
    self.reward_range = [0, 1]
    self.action2lp = []
    for lesson in range(self.LESSON_SIZE):
      for pos in range(self.CLASS_SIZE * self.WEEK_SIZE * self.DAY_SIZE):
        self.action2lp.append((lesson + 1, pos))
    self.reset()

  def reset(self):
    """
    状態を初期化し、初期の観測値を返す
    """
    self.damage = 0
    self.total = 0
    self.progress = 0
    self.table = self.TIME_TABLE.copy()
    self.lesson = list(self.LESSON_LIST)
    return self.table.copy()

  def step(self, action):
    """
   1ステップ進める処理を記述。戻り値は observation, reward, done(ゲーム終了したか), info(追加の情報の辞書)
    """
    # 配置する授業を取り出す
    lesson, pos = self.action2lp[action]

    # 授業を配置することが可能化どうか
    if self._is_bookable(pos, lesson):
      table = self.table.copy()
      shape = table.shape
      table = table.reshape(1, shape[0] * shape[1])
      if table[0][pos] == lesson:
        table[0][pos] = 0
        self.table = table.reshape(shape)
        self.progress -= 1
        moved = None
      else:
        table[0][pos] = lesson
        self.table = table.reshape(shape)
        self.progress += 1
        moved = True
      self.damage = 0
    else:
      moved = False

    # いまの状態
    observation = self.table.copy()
    # 今回の移動によるダメージ
    self.damage += self._get_damage(moved)
    # 今回の移動による報酬
    reward = self._get_reward(moved)
    self.total += reward
    # 終了するか継続するかの判定
    self.done = self._is_done()
    return observation, reward, self.done, {}
  def _close(self):
    """
    [オプション]環境を閉じて後処理をする
    """
    pass
  def _seed(self, seed=None):
    """
    [オプション]ランダムシードを固定する
    """
    pass
  def _get_reward(self, moved):
    if moved is None:
      return -10
    if moved:
      # 特典の計算。授業が分散されて配置できていると高得点になるようにしたいけど効率的な判定方法分からないので一旦スルー
      # return self.progress
      return 1
    else:
      return 0
  def _get_damage(self, moved):
    if moved or moved is None:
      return 0
    else:
      return 1

  def _is_bookable(self, pos, lesson):
    """
    授業を配置することが可能か否かを返す
    """
    table = self.table.copy()
    shape = self.table.shape
    table = table.reshape(1, shape[0] * shape[1])
    # 配置不可の枠には配置できない
    if table[0][pos] == -1:
      return False
    # すでに配置済みのところには配置できない
    if table[0][pos] != lesson and table[0][pos] > 0:
      return False
    # 更新をしてみる
    if table[0][pos] == lesson:
      return True
    table[0][pos] = lesson
    table = table.reshape(shape)
    tableT = table.T
    # 同じ時間帯に授業が入っていないかを確認する
    for i in range(shape[1]):
      row = tableT[i]
      count = Counter(row)
      if len([num for num in count.items() if num[0] > 0 and num[1] > 1]) > 0:
        return False
    # 1日に同じ授業は2回実施できない
    table = table.reshape(self.CLASS_SIZE, self.WEEK_SIZE, self.DAY_SIZE)
    for i in range(self.CLASS_SIZE):
      for j in range(self.WEEK_SIZE):
        row = table[i][j]
        count = Counter(row)
        if len([num for num in count.items() if num[0] > 0 and num[1] > 1]) > 0:
          return False
    # 週に指定の回数以上の授業を担当のクラス以外に実施できない
    table = table.reshape(self.CLASS_SIZE, self.WEEK_SIZE * self.DAY_SIZE)
    for i in range(self.CLASS_SIZE):
      row = table[i]
      count = Counter(row)
      # 担当のクラス以外の授業は実施できない
      if len([num for num in count.items() if num[0] > 0 and i not in self.CLASS_LESSON[num[0]]]) > 0:
        return False
      # 週に指定の回数以上の授業を実施できない
      if len([num for num in count.items() if num[0] > 0 and num[1] > self.CLASS_LESSON_TIMES[num[0]]]) > 0:
        return False
    return True
  def _is_done(self):
    """
    すべての授業を配置するか、授業を配置できない状態がMAX_DAMAGE回続くと終了
    """
    if self.damage >= self.MAX_DAMAGE:
      return True
    elif self.TOTAL_LESSON - self.progress == 0:
      self.render()
      sys.exit(0)
      return True
    return False

  def render(self, mode='human', close=False):
    """
    環境を可視化する
    """
    # human の場合はコンソールに出力。ansiの場合は StringIO を返す
    outfile = StringIO() if mode == 'ansi' else sys.stdout
    outfile.write('\n'.join(' '.join(self.LESSON_NAME[int(elem)] for elem in row) for row in self.table) + '\n' + \
      'total:' + str(self.total) + '\n' + \
      'damage:' + str(self.damage) + '\n' + \
      'lesson:' + str(self.TOTAL_LESSON - self.progress) + '\n')
    return outfile


# --------------------------------------------------------
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import torch
import torch.nn.functional as F
import numpy as np
from core.utils import *

from core import networks
from core.traj import TrajAgent
from core.loss import *

class DDPG(TrajAgent):
    def __init__(self, num_inputs, action_space, args):
        super(DDPG, self).__init__(num_inputs, action_space, args, name='DDPG')
        self.action_dim = self.action_space.shape[0]
        self.critic_value_dim = 0 if self.sa_channel_concat else self.action_dim
        self.value_loss_func = huber_loss
        self.critic_num_input = self.num_inputs
        self.critic, self.critic_optim, self.critic_scheduler, self.critic_target = get_critic(self)

    def load_weight(self, weights):
        self.policy.load_state_dict(weights[0])
        self.critic.load_state_dict(weights[1])
        self.traj_feature_extractor.load_state_dict(weights[2])
        self.state_feature_extractor.load_state_dict(weights[3])

    def get_weight(self):
        return [self.policy.state_dict(),
                self.critic.state_dict(),
                self.traj_feature_extractor.state_dict(),
                self.state_feature_extractor.state_dict() ]

    def get_mix_ratio(self, update_step):
        """
        Get a mixed schedule for supervised learning and RL
        """
        idx = int((self.update_step > np.array(self.mix_milestones)).sum())
        mix_policy_ratio = get_valid_index(self.mix_policy_ratio_list, idx)
        mix_policy_ratio = min(mix_policy_ratio, self.ddpg_coefficients[4])
        mix_value_ratio  = get_valid_index(self.mix_value_ratio_list, idx)
        mix_value_ratio  = min(mix_value_ratio, self.ddpg_coefficients[3])
        return mix_value_ratio, mix_policy_ratio

    def extract_feature(self, image_batch, point_batch,
                            action_batch=None,  goal_batch=None, time_batch=None,
                            vis=False, value=False, repeat=False,
                            train=True, separate=False, curr_joint=None, traj_latent=None,
                            traj_point_state=None, traj_integer_time=None,  traj_action=None,
                            joint_traj_state=None, next_state=False, traj_time_batch=None):
        """
        extract features for policy learning
        """
        if curr_joint is not None: curr_joint = curr_joint[:, :7]
        feature_2  = value and not self.shared_feature
        use_traj   = False
        time_input   = traj_integer_time  if use_traj else  time_batch
        joint_input  = joint_traj_state   if use_traj else  curr_joint
        point_input  = traj_point_state   if use_traj else  point_batch[:,:,:1030]
        action_input = traj_action        if use_traj else  action_batch

        # clip time
        time_input = torch.clamp(time_input, 0., 25.)

        if self.sa_channel_concat and value:
            point_input = concat_state_action_channelwise(point_input, action_input)

        feature = self.unpack_batch(image_batch, point_input, gt_goal=goal_batch, val=feature_2, repeat=repeat)
        if self.use_time:
            feature[0] = torch.cat((feature[0], time_input[:,None]), dim=1)

        if not self.sa_channel_concat and value:
            feature[0] = torch.cat((feature[0], action_input), dim=1)
        self.state_policy_feat = feature[0].detach().cpu().numpy()
        return feature

    def target_value(self):
        """
        compute target value
        """
        next_time_batch = self.time_batch - 1
        next_traj_time_batch = self.traj_integer_time_batch - 1
        reward_batch = self.traj_reward_batch if self.full_traj_embedding else self.reward_batch
        mask_batch = self.traj_mask_batch if self.full_traj_embedding else self.mask_batch

        with torch.no_grad():
            next_state_batch, _  = self.extract_feature(self.next_image_state_batch,
                                                        self.next_point_state_batch,
                                                        action_batch=self.next_action_batch,
                                                        goal_batch=self.next_goal_batch,
                                                        traj_point_state=self.next_traj_point_state_batch,
                                                        time_batch=next_time_batch,
                                                        curr_joint=self.next_joint_batch,
                                                        joint_traj_state=self.traj_joint_batch,
                                                        traj_integer_time=next_traj_time_batch,
                                                        traj_action=self.next_traj_action_batch,
                                                        separate=True,
                                                        vis=False,
                                                        next_state=True,
                                                        value=False)

            next_action_mean, _, _, _ = self.policy_target.sample(next_state_batch)
            idx = int((self.update_step > np.array(self.mix_milestones)).sum())
            noise_scale = self.action_noise * get_valid_index(self.noise_ratio_list, idx)
            noise_delta = get_noise_delta(next_action_mean, noise_scale, self.noise_type)
            noise_delta[:, :3] = torch.clamp(noise_delta[:, :3], -0.01, 0.01)
            next_action_mean = next_action_mean + noise_delta
            if self.sa_channel_concat or not self.shared_feature:
                next_target_state_batch, _ = self.extract_feature( self.next_image_state_batch, self.next_point_state_batch, next_action_mean,
                                                                   self.next_goal_batch, next_time_batch, value=True, curr_joint=self.next_joint_batch)
            else:
                next_target_state_batch = torch.cat((next_state_batch, next_action_mean), dim=-1)

            min_qf_next_target = self.state_action_value(next_target_state_batch, next_action_mean, target=True)
            next_q_value = reward_batch + (1 - mask_batch) * self.gamma * min_qf_next_target
            return next_q_value

    def state_action_value(self, state_batch, action_batch, target=False, return_qf=False):
        """ compute state action value """
        critic = self.critic_target if target else self.critic
        qf1, qf2, critic_aux = critic(state_batch, None)
        min_qf = torch.min(qf1, qf2)
        min_qf = min_qf.squeeze(-1)

        if return_qf:
            return qf1.squeeze(-1), qf2.squeeze(-1), critic_aux
        return min_qf

    def critic_optimize(self):
        """
        optimize critic and feature gradient
        """
        self.critic_optim.zero_grad()
        self.state_feat_val_encoder_optim.zero_grad()
        critic_loss = sum([getattr(self, name) for name in get_loss_info_dict().keys() \
                            if name.endswith('loss') and name.startswith('critic')])
        if self.sa_channel_concat or not self.shared_feature:
            critic_loss.backward()
        else:
            critic_loss.backward(retain_graph=True)
        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), self.clip_grad)
        if self.train_value_feature: self.state_feat_val_encoder_optim.step()
        self.critic_optim.step()

    def compute_critic_loss(self, value_feat):
        """
        compute one step bellman error
        """
        self.next_q_value = self.target_value()
        self.qf1, self.qf2, self.critic_grasp_aux = self.state_action_value(value_feat, self.action_batch, return_qf=True)
        self.critic_loss = F.mse_loss(self.qf1.view(-1), self.next_q_value) + \
                           F.mse_loss(self.qf2.view(-1), self.next_q_value)

        if self.critic_aux:
            self.critic_grasp_aux_loss += goal_pred_loss(self.critic_grasp_aux[self.goal_reward_mask, :7], self.goal_batch[self.goal_reward_mask])

    def update_parameters(self, batch_data, updates, k, test=False):
        """ update agent parameters """
        self.mix_value_ratio, self.mix_policy_ratio = self.get_mix_ratio(self.update_step)
        self.set_mode(test)
        self.prepare_data(batch_data)

        if self.train_traj_feature :
            self.train_traj_net()

        # value
        if self.policy_update_gap > 0:
            value_feat, _ = self.extract_feature(self.image_state_batch,
                                                self.point_state_batch,
                                                action_batch=self.action_batch,
                                                goal_batch=self.goal_batch,
                                                traj_point_state=self.traj_point_state_batch,
                                                time_batch=self.time_batch,
                                                curr_joint=self.curr_joint,
                                                joint_traj_state=self.traj_joint_batch,
                                                traj_integer_time=self.traj_integer_time_batch,
                                                traj_time_batch=self.traj_time_batch,
                                                traj_action=self.traj_action_batch,
                                                value=True)
            self.compute_critic_loss(value_feat)
            self.critic_optimize()

        # policy
        if self.sa_channel_concat or not self.shared_feature:
            policy_feat, extra_feat_pred = self.extract_feature(self.image_state_batch,
                                                            self.point_state_batch,
                                                            goal_batch=self.goal_batch,
                                                            traj_point_state=self.traj_point_state_batch,
                                                            time_batch=self.time_batch,
                                                            curr_joint=self.curr_joint,
                                                            separate=True,
                                                            joint_traj_state=self.traj_joint_batch,
                                                            traj_integer_time=self.traj_integer_time_batch,
                                                            traj_time_batch=self.traj_time_batch)
        else:
            policy_feat = value_feat[..., :-self.action_dim]

        self.pi, _, _, self.aux_pred = self.policy.sample(policy_feat)

        # actor critic
        if self.has_critic and (self.update_step % self.policy_update_gap == 0) and self.policy_update_gap > 0:
            if self.sa_channel_concat or not self.shared_feature:
                value_pi_feat, _  = self.extract_feature(self.image_state_batch, self.point_state_batch,
                                                         self.pi, self.goal_batch, self.time_batch,
                                                         value=True, curr_joint=self.curr_joint)
            else:
                value_pi_feat = torch.cat((policy_feat, self.pi), dim=-1)
            self.qf1_pi, self.qf2_pi, critic_aux = self.state_action_value(value_pi_feat, self.pi, return_qf=True)
            self.qf1_pi = self.qf1_pi[~self.expert_reward_mask]
            self.qf2_pi = self.qf2_pi[~self.expert_reward_mask]
            self.actor_critic_loss = -self.mix_policy_ratio * self.qf1_pi.mean()

        loss = self.compute_loss()
        self.optimize(loss, self.update_step)
        self.update_step += 1

        # log
        self.log_stat()
        return {k: float(getattr(self, k)) for k in get_loss_info_dict().keys()}


# coding: utf-8

# This file builds a object detection tensorflow using tensorflow api and
# the process of building the deep neural net is derived from object_detection_tutorial.ipynb

# Needed libraries
from utils import visualization_utils as vis_util
from utils import label_map_util
from matplotlib import pyplot as plt
import six.moves.urllib as urllib
import tensorflow as tf
from io import StringIO
from PIL import Image
import numpy as np
import tarfile
import zipfile
import sys
import cv2
import os

class objectDetection:
  def __init__(self, modelPath, labelMapPath):
    self.modelPath = modelPath
    self.labelMapPath = labelMapPath

  # model setup:
  # 1. Specify which tensorflow model to use
  # 2. Locate label files and indicate how many classes in the labels
  # 3. Load in label map using the label file
  # 4. Return the category indices which correspond to the digital labels
  def model_setup(self):
    PATH_TO_CKPT = ""
    PATH_TO_LABELS = ""
    # if model is not specified, use default ssd_mobilenet_v1_coco
    if self.modelPath == None or self.labelMapPath== "":
      # What model to download.
      MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017'
      MODEL_FILE = MODEL_NAME + '.tar.gz'
      DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'
      # Path to frozen detection graph. This is the actual model that is used for the object detection.
      PATH_TO_CKPT = "./submodule/" + MODEL_NAME + '/frozen_inference_graph.pb'
      if os.path.isfile(PATH_TO_CKPT) and os.access(PATH_TO_CKPT, os.R_OK):
        pass
      else:
        # Download Model
        opener = urllib.request.URLopener()
        opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
        tar_file = tarfile.open(MODEL_FILE)
        for file in tar_file.getmembers():
          file_name = os.path.basename(file.name)
          if 'frozen_inference_graph.pb' in file_name:
            tar_file.extract(file, os.getcwd() + "/submodule")
    else:
      PATH_TO_CKPT = self.modelPath

    if self.labelMapPath == None or self.labelMapPath == "":
      # List of the strings that is used to add correct label for each box.
      PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')
    else:
      PATH_TO_LABELS = self.labelMapPath

    NUM_CLASSES = 90

    # Load a (frozen) Tensorflow model into memory.
    # Use pre-trained model
    global detection_graph
    detection_graph = tf.Graph()
    with detection_graph.as_default():
      od_graph_def = tf.GraphDef()
      with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

    # Loading label map
    # Label maps map indices to category names, so that when our convolution network predicts `5`,
    # we know that this corresponds to `airplane`.  Here we use internal utility functions,
    # but anything that returns a dictionary mapping integers to appropriate string labels would be fine
    label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
    categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
    category_index = label_map_util.create_category_index(categories)
    return category_index

  # Any model exported using the `export_inference_graph.py` tool can be loaded here simply
  # by changing `PATH_TO_CKPT` to point to a new .pb file.

  # By default we use an "SSD with Mobilenet" model here.
  # See the [detection model zoo](https://github.com/tensorflow/models/blob/master/research/ \
  # object_detection/g3doc/detection_model_zoo.md) for a list of other models
  # that can be run out-of-the-box with varying speeds and accuracies.

  # detection process:
  # Startup tensorflow process and draw boxes around detected objects in the pictures/videos/camera
  # stream with predicted name of the detected object and percentage of correctness confidence
  # Argument: cap refers to the input to be detected (can be images, videos, or streaming camera)
  #           flag refers to a boolean indicating whether the input is a static image or not
  # Return: True if the video/camera detection process is interrupted by user or lost of streaming input.
  # Keep detecting if input stream continues
  def detect_process(self, cap, flag):
    category_index = self.model_setup()
    with detection_graph.as_default():
      with tf.Session(graph=detection_graph) as sess:
        # Definite input and output Tensors for detection_graph
        image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
        detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
        num_detections = detection_graph.get_tensor_by_name('num_detections:0')

        # detection process loop
        while True:
          # input is video or camera
          if flag == False:
            ret, frame = cap.read()
            if ret == True:
              image_np = frame
            # If input stream is lost,
            # terminate the process by return the last frame that is after detection
            else:
              flag = True
          # input is a static image
          else:
              image_np = cap
          # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
          image_np_expanded = np.expand_dims(image_np, axis=0)
          # Actual detection.
          (boxes, scores, classes, num) = sess.run(
              [detection_boxes, detection_scores, detection_classes, num_detections],
              feed_dict={image_tensor: image_np_expanded})
          # Visualization of the results of a detection.
          image_done = vis_util.visualize_boxes_and_labels_on_image_array(
              image_np,
              np.squeeze(boxes),
              np.squeeze(classes).astype(np.int32),
              np.squeeze(scores),
              category_index,
              use_normalized_coordinates=True,
              line_thickness=8)
          # if input is an image, just return the detected version once it comes
          # back from the tensorflow graph
          if flag == True:
            output_img = cv2.resize(image_done, (1024, 768))
            return output_img
          # if input is an image, just return the detected version once it comes
          # back from the tensorflow graph
          else:
            cv2.imshow('object detection', cv2.resize(image_done, (1024, 768)))
            # Quit detection if user interrupts it by pressing the key 'q'
            if cv2.waitKey(25) & 0xFF == ord('q'):
              cv2.destroyAllWindows()
              break
        # return a signal of process termination
        return True