# Copyright 2021-2022 Xanadu Quantum Technologies Inc.

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

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

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module contains a function to run aa custom PennyLane VQE problem on qiskit runtime.
"""
# pylint: disable=too-few-public-methods,protected-access,too-many-arguments,too-many-branches,too-many-statements

import os
import warnings
import inspect
from collections import OrderedDict

import pennylane.numpy as np
import pennylane as qml

from pennylane_qiskit.qiskit_device import QiskitDevice
from pennylane_qiskit.ibmq import connect
import qiskit.circuit.library.n_local as lib_local
from qiskit.providers.ibmq.runtime import ResultDecoder
from qiskit.circuit import ParameterVector, QuantumCircuit, QuantumRegister
from qiskit.converters import circuit_to_dag, dag_to_circuit
from qiskit import IBMQ

from scipy.optimize import OptimizeResult

class VQEResultDecoder(ResultDecoder):
    """The class is used to decode the result from the runtime problem and return it as a
    Scipy Optimizer result.
    """

    @classmethod
    def decode(cls, data):
        """Decode the data from the VQE program."""
        data = super().decode(data)
        return OptimizeResult(data)

class RuntimeJobWrapper:
    """A simple Job wrapper that attaches intermediate results directly to the job object itself
    in the `intermediate_results attribute` via the `_callback` function.
    """

    def __init__(self):
        self._job = None
        self._decoder = VQEResultDecoder
        self.intermediate_results = {
            "nfev": [],
            "parameters": [],
            "function": [],
            "step": [],
            "accepted": [],
        }

    def _callback(self, *args):
        """The callback function that attaches intermediate results to the wrapper:

        Args:
            nfev (int): Number of evaluation.
            xk (array_like): A list or NumPy array to attach.
            fk (float): Value of the function.
            step (float): Value of the step.
            accepted (bool): True if the loss function value has improved, False otherwise.
        """
        # If it is a dictionary it is the final result and does not belong to intermediate results
        if not isinstance(args[1], dict):
            _, (nfev, xk, fk, step, accepted) = args
            self.intermediate_results["nfev"].append(nfev)
            self.intermediate_results["parameters"].append(xk)
            self.intermediate_results["function"].append(fk)
            self.intermediate_results["step"].append(step)
            self.intermediate_results["accepted"].append(accepted)

    def _scipy_callback(self, *args):
        """The callback function that attaches intermediate results to the wrapper:

        Args:
            xk (array_like): A list or NumPy array to attach.
        """
        # If it is a dictionary it is the final result and does not belong to intermediate results
        if not isinstance(args[1], dict):
            _, xk = args
            self.intermediate_results["parameters"].append(xk)

    def result(self):
        """Get the result of the job as a SciPy OptimizerResult object.

        This method blocks until the job is done, cancelled, or raises an error.

        Returns:
            OptimizerResult: An optimizer result object.
        """
        return self._job.result(decoder=self._decoder)

def upload_vqe_runner(hub="ibm-q", group="open", project="main", **kwargs):
    """Upload the custom VQE runner to the IBMQ cloud.

    Args:
        hub (str): Ibmq provider hub.
        group (str): Ibmq provider group.
        project (str): Ibmq provider project.

    Returns:
        str: Program id that can be used to run the program.
    """

    connect(kwargs)

    # Specify a single hub, group and project
    hub = kwargs.get("hub", "ibm-q")
    group = kwargs.get("group", "open")
    project = kwargs.get("project", "main")

    meta = {
        "name": "vqe-runtime",
        "description": "A sample VQE program.",
        "max_execution_time": 100000,
        "spec": {},
    }

    meta["spec"]["parameters"] = {
        "$schema": "https://json-schema.org/draft/2019-09/schema",
        "properties": {
            "hamiltonian": {
                "description": "Hamiltonian whose ground state we want to find.",
                "type": "array",
            },
            "x0": {
                "description": "Initial vector of parameters for the quantum circuit.",
                "type": "array",
            },
            "ansatz": {
                "description": "Qiskit circuit or name of ansatz quantum circuit to use, default='EfficientSU2'",
                "type": "[QuantumCircuit,string]",
                "default": "EfficientSU2",
            },
            "ansatz_config": {
                "description": "Configuration parameters for the ansatz circuit.",
                "type": "dict",
            },
            "optimizer": {
                "description": "Classical optimizer to use, default='SPSA'.",
                "type": "string",
                "default": "SPSA",
            },
            "optimizer_config": {
                "description": "Configuration parameters for the optimizer.",
                "type": "dict",
            },
            "shots": {
                "description": "The number of shots used for each circuit evaluation.",
                "type": "integer",
            },
            "use_measurement_mitigation": {
                "description": "Use measurement mitigation, default=False.",
                "type": "boolean",
                "default": False,
            },
        },
        "required": ["hamiltonian", "x0"],
    }

    meta["spec"]["return_values"] = {
        "$schema": "https://json-schema.org/draft/2019-09/schema",
        "description": "Final result in SciPy optimizer format",
        "type": "object",
    }

    meta["spec"]["intermadiate_results"] = {
        "$schema": "https://json-schema.org/draft/2019-09/schema",
        "description": "Dictionnary containing: "
        "The number of evaluation at current optimization step."
        "Parameter vector at current optimization step."
        "Function value at the current optimization step."
        "The size of the step.",
        "type": "dict",
    }

    provider = IBMQ.get_provider(hub=hub, group=group, project=project)
    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
    PROG_FILE = "vqe_runtime_program.py"
    program_path = os.path.join(ROOT_DIR, PROG_FILE)

    program_id = provider.runtime.upload_program(data=program_path, metadata=meta)
    return program_id

def delete_vqe_runner(provider, program_id):
    """Delete the desired program on the IBMQ platform.
    Args:
        provider (object): IBMQ provider.
        program_id (str): Id of the Qiskit runtime to be deleted.
    """
    provider.runtime.delete_program(program_id)

def vqe_runner(
    backend,
    hamiltonian,
    x0,
    program_id,
    ansatz="EfficientSU2",
    ansatz_config=None,
    optimizer="SPSA",
    optimizer_config=None,
    shots=8192,
    use_measurement_mitigation=False,
    **kwargs,
):
    """Routine that executes a given VQE problem via the sample-vqe program on the target backend.

    Args:
        backend (ProgramBackend): Qiskit backend instance.
        hamiltonian (qml.Hamiltonian): Hamiltonian whose ground state we want to find.
        x0 (array_like): Initial vector of parameters.
        program_id(str): Id of the program, it has to be generated by using the upload_vqe_runner function.
        Once the program is uploaded, you can find the id in your program list online.
        ansatz (Quantum function or str): Optional, a PennyLane quantum function or the name of the Qiskit
            ansatz quantum circuit to use. Default='EfficientSU2'
        ansatz_config (dict): Optional, configuration parameters for the ansatz circuit if from Qiskit library.
        optimizer (str): Optional, string specifying classical optimizer. Default='SPSA'.
        optimizer_config (dict): Optional, configuration parameters for the optimizer.
        shots (int): Optional, number of shots to take per circuit. Default=1024.
        use_measurement_mitigation (bool): Optional, use measurement mitigation. Default=False.

    Returns:
        OptimizeResult: The result in SciPy optimization format.
    """
    # Init the dictionnaries
    if ansatz_config is None:
        ansatz_config = {}

    if optimizer_config is None:
        optimizer_config = {"maxiter": 100}

    if not isinstance(hamiltonian, qml.Hamiltonian):
        raise qml.QuantumFunctionError("A PennyLane Hamiltonian object is required.")

    connect(kwargs)

    options = {"backend_name": backend}

    inputs = {}

    # Validate circuit ansatz and number of qubits
    if not isinstance(ansatz, str):
        inputs["x0"], inputs["ansatz"], num_qubits, wires = _pennylane_to_qiskit_ansatz(
            ansatz, x0, hamiltonian
        )

    # The circuit will be taken from the Qiskit library as a str was passed.
    else:
        wires = hamiltonian.wires
        num_qubits = len(wires)

        ansatz_circ = getattr(lib_local, ansatz, None)
        if ansatz_circ is None:
            raise ValueError(f"Ansatz {ansatz} not in n_local circuit library.")

        inputs["ansatz"] = ansatz
        inputs["ansatz_config"] = ansatz_config

        # If given x0, validate its length against num_params in ansatz
        x0 = np.asarray(x0)
        ansatz_circ = ansatz_circ(num_qubits, **ansatz_config)
        num_params = ansatz_circ.num_parameters

        if x0.shape[0] != num_params:
            warnings.warn(
                "The shape of parameters array is not correct, a random initialization has been applied."
            )
            x0 = 2 * np.pi * np.random.rand(num_params)

        inputs["x0"] = x0

    # Transform the PennyLane hamilonian to a suitable form
    hamiltonian = hamiltonian_to_list_string(hamiltonian, wires)

    inputs["hamiltonian"] = hamiltonian

    # Set the rest of the inputs
    inputs["optimizer"] = optimizer
    inputs["optimizer_config"] = optimizer_config
    inputs["shots"] = shots
    inputs["use_measurement_mitigation"] = use_measurement_mitigation

    # Specify a single hub, group and project
    hub = kwargs.get("hub", "ibm-q")
    group = kwargs.get("group", "open")
    project = kwargs.get("project", "main")

    provider = IBMQ.get_provider(hub=hub, group=group, project=project)

    rt_job = RuntimeJobWrapper()

    # Callbacks functions are different between optimizers.
    if optimizer in ["SPSA", "QNSPSA"]:
        job = provider.runtime.run(
            program_id, options=options, inputs=inputs, callback=rt_job._callback
        )
    else:
        job = provider.runtime.run(
            program_id, options=options, inputs=inputs, callback=rt_job._scipy_callback
        )
    rt_job._job = job

    return rt_job

def _pennylane_to_qiskit_ansatz(ansatz, x0, hamiltonian):
    r"""Convert an ansatz from PennyLane to a circuit in Qiskit.

    Args:
        ansatz (Quantum Function): A PennyLane quantum function that represents the circuit.
        x0 (array_like): The array of parameters.
        num_qubits_h (int): Number of qubits evaluated from the Hamiltonian.

    Returns:
        list[tuple[float,str]]: Hamiltonian in a format for the runtime program.
    """

    if isinstance(ansatz, (qml.QNode, qml.tape.QuantumTape)):
        raise qml.QuantumFunctionError("The ansatz must be a callable quantum function.")

    if callable(ansatz):
        if len(inspect.getfullargspec(ansatz).args) != 1:
            raise qml.QuantumFunctionError("Param should be a single vector.")
        try:
            tape_param = x0[0] if len(x0) == 1 else x0
            tape = qml.transforms.make_tape(ansatz)(np.array(tape_param)).expand(
                depth=5, stop_at=lambda obj: obj.name in QiskitDevice._operation_map
            )
        except IndexError as e:
            raise qml.QuantumFunctionError("Not enough parameters in X0.") from e

        # Raise exception if there are no operations
        if len(tape.operations) == 0:
            raise qml.QuantumFunctionError("Function contains no quantum operations.")

        params = tape.get_parameters()
        trainable_params = []

        for p in params:
            if qml.math.requires_grad(p):
                trainable_params.append(p)

        num_params = len(trainable_params)

        if len(x0) != num_params:
            warnings.warn(
                "In order to match the tape expansion, the number of parameters has been changed."
            )
            x0 = 2 * np.pi * np.random.rand(num_params)

        wires_circuit = tape.wires
        wires_hamiltonian = hamiltonian.wires
        all_wires = wires_circuit + wires_hamiltonian

        # Set the number of qubits
        num_qubits = len(all_wires)

        circuit_ansatz = _qiskit_ansatz(num_params, num_qubits, all_wires, tape)

    else:
        raise ValueError("Input ansatz is not a quantum function or a string.")

    return x0, circuit_ansatz, num_qubits, all_wires

def _qiskit_ansatz(num_params, num_qubits, wires, tape):
    """Transform a quantum tape from PennyLane to a Qiskit circuit.

    Args:
        num_params (int): Number of parameters.
        num_qubits (int): Number of qubits.
        wires (qml.wire.Wires): Wires used in the tape and Hamiltonian.
        tape (qml.tape.QuantumTape): The quantum tape of the circuit ansatz in PennyLane.

    Returns:
        QuantumCircuit: Qiskit quantum circuit.

    """
    consecutive_wires = qml.wires.Wires(range(num_qubits))
    wires_map = OrderedDict(zip(wires, consecutive_wires))
    # From here: Create the Qiskit ansatz circuit
    params_vector = ParameterVector("p", num_params)

    reg = QuantumRegister(num_qubits, "q")
    circuit_ansatz = QuantumCircuit(reg, name="vqe")

    circuits = []

    j = 0
    for operation in tape.operations:
        wires = operation.wires.map(wires_map)
        par = operation.parameters
        operation = operation.name
        mapped_operation = QiskitDevice._operation_map[operation]

        qregs = [reg[i] for i in wires.labels]

        if operation.split(".inv")[0] in ("QubitUnitary", "QubitStateVector"):
            # Need to revert the order of the quantum registers used in
            # Qiskit such that it matches the PennyLane ordering
            qregs = list(reversed(qregs))

        dag = circuit_to_dag(QuantumCircuit(reg, name=""))

        if operation in ("QubitUnitary", "QubitStateVector"):
            # Parameters are matrices
            gate = mapped_operation(par[0])
        else:
            # Parameters for the operation
            if par and qml.math.requires_grad(par[0]):
                op_num_params = len(par)
                par = []
                for num in range(op_num_params):
                    par.append(params_vector[j + num])
                j += op_num_params

            gate = mapped_operation(*par)

        if operation.endswith(".inv"):
            gate = gate.inverse()

        dag.apply_operation_back(gate, qargs=qregs)
        circuit = dag_to_circuit(dag)
        circuits.append(circuit)

    for circuit in circuits:
        circuit_ansatz &= circuit

    return circuit_ansatz

def hamiltonian_to_list_string(hamiltonian, wires):
    r"""Convert a Hamiltonian object from PennyLane to a list of pairs representing each coefficient and
    term in the Hamiltonian.

    Args:
        hamiltonian (qml.Hamiltonian): A Hamiltonian from PennyLane.
        wires (qml.wires.Wires): A list of qubits from PennyLane.

    Returns:
        list[tuple[float,str]]: Hamiltonian in a format for the runtime program.
    """

    num_qubits = len(wires)

    consecutive_wires = qml.wires.Wires(range(num_qubits))
    wires_map = OrderedDict(zip(wires, consecutive_wires))

    coeff, observables = hamiltonian.terms()

    authorized_obs = {"PauliX", "PauliY", "PauliZ", "Hadamard", "Identity"}

    for obs in observables:
        obs_names = obs.name if isinstance(obs.name, list) else [obs.name]
        if any(ob not in authorized_obs for ob in obs_names):
            raise qml.QuantumFunctionError("Observable is not accepted.")

    # Create string Hamiltonian
    obs_str = {"PauliX": "X", "PauliY": "Y", "PauliZ": "Z", "Hadamard": "H", "Identity": "I"}

    obs_org = []
    # Map the PennyLane hamiltonian to a list PauliY(1) @ PauliY(0) -> [[[0,'Y'], [1,'Y']]]
    for obs in observables:
        # Tensors
        if isinstance(obs.name, list):
            internal = []
            for i, j in zip(obs.wires.map(wires_map).tolist(), obs.name):
                internal.append([i, obs_str[j]])
            internal.sort()
            obs_org.append(internal)
        else:
            obs_org.append([[obs.wires.map(wires_map).tolist()[0], obs_str[obs.name]]])

    # Create the hamiltonian terms as lists of strings [[[0,'Y'], [1,'Y']]] -> [['YI'], ['IY']]
    obs_list = []
    for elem in obs_org:
        empty_obs = ["I"] * num_qubits
        for el in elem:
            wire = el[0]
            observable = el[1]
            empty_obs[wire] = observable
        obs_list.append(empty_obs)

    # Create the list of tuples with coeffs and Hamiltonian terms as strings [['YI'], ['IY']] -> [(1, 'YI'), (1, 'IY')]
    hamiltonian = []
    for i, elem in enumerate(obs_list):
        result = "".join(elem)
        hamiltonian.append((coeff[i], result))

    return hamiltonian


import sys,string,os,shutil
sys.path.insert(1, os.path.join(sys.path[0], '..')) ### import parent dir dependencies
from scipy import sparse, io
import numpy
import LineageProfilerIterate
import cluster_corr
import export
from import_scripts import ChromiumProcessing
import traceback

""" cellHarmony without alignment """

def cellHarmony(species,platform,query_exp_file,exp_output,
            customMarkers=False,useMulti=False,fl=None,customLabels=None):
    """ Prepare pre-aligned result files in a pre-defined format for cellHarmony post-aligment
    differential and visualization analyses """

    customLabels = fl.Labels()
    reference_exp_file = customMarkers ### pre-formatted from Seurat or other outputs

    export_directory = os.path.abspath(os.path.join(query_exp_file, os.pardir))
    if 'ExpressionInput' in query_exp_file:
        ### Change to the root directory above ExpressionINput
        export_directory = os.path.abspath(os.path.join(export_directory, os.pardir))
    dataset_name = string.replace(string.split(query_exp_file,'/')[-1][:-4],'exp.','')
    try: os.mkdir(export_directory+'/cellHarmony/')
    except: pass
    try: os.mkdir(export_directory+'/cellHarmony/CellClassification/')
    except: pass
    try: os.mkdir(export_directory+'/cellHarmony/OtherFiles/')
    except: pass

    ### Get the query and reference cells, dataset names
    refererence_cells, query_cells, reference_dataset, query_dataset = importCelltoClusterAnnotations(customLabels)  ### Get the reference and query cells in their respective order

    ### copy and re-name the input expression file to the output cellHarmony directory
    if len(reference_dataset)>0 and len(query_dataset)>0:
        target_exp_dir = export_directory+'/cellHarmony/exp.'+reference_dataset+'__'+query_dataset+'-AllCells.txt'
    else:
        target_exp_dir = export_directory+'/cellHarmony/exp.cellHarmony-reference__Query-AllCells.txt'
        reference_dataset = 'cellHarmony-reference'
    shutil.copy(query_exp_file,target_exp_dir)

    ### filter and export the heatmap with just reference cells
    cell_cluster_order = simpleHeaderImport(reference_exp_file)
    filtered_reference_cells=[]
    filtered_query_cells_db={}
    filtered_query_cells=[]
    representative_refcluster_cell = {}
    for cell_id in cell_cluster_order:
        if cell_id in refererence_cells:
            filtered_reference_cells.append(cell_id)
            cluster_label = refererence_cells[cell_id].Label()
            ### Identifies where to place each query cell
            try: representative_refcluster_cell[cluster_label].append(cell_id)
            except: representative_refcluster_cell[cluster_label] = [cell_id]
        elif cell_id in query_cells:
            filtered_query_cells_db[cell_id]=query_cells[cell_id]
            filtered_query_cells.append(cell_id)

    #reference_output_file = export.findParentDir(reference_exp_file)+'/'+reference_dataset+'.txt'
    reference_output_file = export_directory+'/cellHarmony/OtherFiles/'+reference_dataset+'.txt'
    reference_output_file2 = export_directory+'/cellHarmony/exp.'+reference_dataset+'__'+query_dataset+'-Reference.txt'
    query_output_file =export_directory+'/'+query_dataset+'.txt'
    ### Write out separate refernece and query files
    from import_scripts import sampleIndexSelection
    sampleIndexSelection.filterFile(reference_exp_file,reference_output_file,['row_clusters-flat']+filtered_reference_cells,force=True)
    sampleIndexSelection.filterFile(target_exp_dir,query_output_file,filtered_query_cells,force=True)
    shutil.copy(reference_output_file,reference_output_file2)

    ### export the CellClassification file
    output_classification_file = export_directory+'/cellHarmony/CellClassification/CellClassification.txt'
    exportCellClassifications(output_classification_file,filtered_query_cells_db,filtered_query_cells,representative_refcluster_cell)
    labels_file = export_directory+'/labels.txt'
    exportLabels(labels_file,filtered_reference_cells,refererence_cells)
    fl.setLabels(labels_file)

    print 'Files formatted for cellHarmony... running differential expression analyses'
    try:
        print reference_output_file
        print query_output_file
        print output_classification_file
        LineageProfilerIterate.harmonizeClassifiedSamples(species, reference_output_file, query_output_file, output_classification_file,fl=fl)
    except:
        print '\nFAILED TO COMPLETE THE FULL CELLHARMONY ANALYSIS (SEE LOG FILE)...'
        print traceback.format_exc()

    return True

def exportCellClassifications(output_file,query_cells,filtered_query_cells,representative_refcluster_cell):
    """ Match the Louvain cellHarmony export format for the classification file """

    header = 'Query Barcode\tRef Barcode\tCorrelation\tQuery Partition\tRef Partition\tLabel\n'

    o = open(output_file,'w')
    o.write(header)
    for query_barcode in filtered_query_cells:
        CI = query_cells[query_barcode]
        cluster_number = CI.ClusterNumber()
        label = CI.Label()
        try:
            ref_barcode = representative_refcluster_cell[label][-1]
        except:
            continue
        values = [query_barcode,ref_barcode,'1.0',cluster_number,cluster_number,label]
        o.write(string.join(values,'\t')+'\n')
    o.close()

def exportLabels(labels_file,filtered_reference_cells,refererence_cells):
    l = open(labels_file,'w')
    for cell_id in filtered_reference_cells:
        CI = refererence_cells[cell_id]
        cluster_number = CI.ClusterNumber()
        label = CI.Label()
        values = [cell_id,cluster_number,label]
        l.write(string.join(values,'\t')+'\n')
    l.close()

def simpleHeaderImport(filename):
    for line in open(filename,'rU').xreadlines():
        data = cleanUpLine(line)
        if '\t' in data:
            t = string.split(data,'\t')
        else:
            t = string.split(data,',')
        header = t[2:]
        header2 = []
        for h in header:
            if ":" in h:
                h = string.split(h,':')[-1]
            header2.append(h)
        break
    return header2

class CellInfo:
    def __init__(self,cell_id, cluster_number, dataset_name, dataset_type, label):
        self.cell_id = cell_id; self.cluster_number = cluster_number; self.dataset_name = dataset_name
        self.dataset_type = dataset_type; self.label = label
    def CellID(self): return self.cell_id
    def ClusterNumber(self): return self.cluster_number
    def DatasetName(self): return self.dataset_name
    def DataSetType(self): return self.dataset_type
    def Label(self): return self.label
    def __repr__(self):
        return self.CellID()+'|'+self.Label()+'|'+self.DataSetType()

def importCelltoClusterAnnotations(filename):
    firstRow = True
    refererence_cells={}
    query_cells={}
    for line in open(filename,'rU').xreadlines():
        data = cleanUpLine(line)
        if '\t' in data:
            t = string.split(data,'\t')
        else:
            t = string.split(data,',')
        if firstRow:
            ci = t.index('cell_id')
            try: cn = t.index('cluster_number')
            except: cn = 'False'
            try: cm = t.index('cluster_name')
            except: cm = 'False'
            try: cnm = t.index('ClustNameNum')
            except: cnm = 'False'
            try: cnm = t.index('label')
            except: pass
            dn = t.index('dataset_name')
            dt = t.index('dataset_type')
            firstRow = False
        else:
            cell_id = t[ci]
            try: cluster_number = t[cn]
            except: cluster_number = 'False'
            dataset_name = t[dn]
            dataset_type = t[dt]
            if cnm !='False':
                label = t[cnm]
                if cluster_number == 'False':
                    cluster_number = label
            elif cm != False:
                cluster_name = t[cm]
                label = cluster_name + '_c'+cluster_number
            elif cluster_number == False:
                label = t[cm]
            else:
                label = 'c'+cluster_number
            if string.lower(dataset_type)[0] == 'r':
                dataset_type = 'Reference'
                reference_dataset = dataset_name
                CI = CellInfo(cell_id, cluster_number, dataset_name, dataset_type, label)
                refererence_cells[cell_id]=CI
            elif string.lower(dataset_type)[0] == 'q':
                dataset_type = 'Query'
                query_dataset = dataset_name
                CI = CellInfo(cell_id, cluster_number, dataset_name, dataset_type, label)
                query_cells[cell_id]=CI
    return refererence_cells, query_cells, reference_dataset, query_dataset

def cleanUpLine(line):
    line = string.replace(line,'\n','')
    line = string.replace(line,'\c','')
    data = string.replace(line,'\r','')
    data = string.replace(data,'"','')
    return data

if __name__ == '__main__':
    platform = 'RNASeq'
    cellHarmony(genome,platform,args.query_h5,None,
            customMarkers=args.reference_h5,useMulti=False,fl=None,customLabels=labels)


from typing import List, Optional, NamedTuple
from mlagents.torch_utils import torch
import numpy as np

from mlagents.trainers.torch.utils import ModelUtils
from mlagents.trainers.buffer import AgentBuffer, BufferKey
from mlagents_envs.base_env import _ActionTupleBase

class LogProbsTuple(_ActionTupleBase):
    """
    An object whose fields correspond to the log probs of actions of different types.
    Continuous and discrete are numpy arrays
    Dimensions are of (n_agents, continuous_size) and (n_agents, discrete_size),
    respectively. Note, this also holds when continuous or discrete size is
    zero.
    """

    @property
    def discrete_dtype(self) -> np.dtype:
        """
        The dtype of a discrete log probability.
        """
        return np.float32

class ActionLogProbs(NamedTuple):
    """
    A NamedTuple containing the tensor for continuous log probs and list of tensors for
    discrete log probs of individual actions as well as all the log probs for an entire branch.
    Utility functions provide numpy <=> tensor conversions to be used by the optimizers.
    :param continuous_tensor: Torch tensor corresponding to log probs of continuous actions
    :param discrete_list: List of Torch tensors each corresponding to log probs of the discrete actions that were
    sampled.
    :param all_discrete_list: List of Torch tensors each corresponding to all log probs of
    a discrete action branch, even the discrete actions that were not sampled. all_discrete_list is a list of Tensors,
    each Tensor corresponds to one discrete branch log probabilities.
    """

    continuous_tensor: torch.Tensor
    discrete_list: Optional[List[torch.Tensor]]
    all_discrete_list: Optional[List[torch.Tensor]]

    @property
    def discrete_tensor(self):
        """
        Returns the discrete log probs list as a stacked tensor
        """
        return torch.stack(self.discrete_list, dim=-1)

    @property
    def all_discrete_tensor(self):
        """
        Returns the discrete log probs of each branch as a tensor
        """
        return torch.cat(self.all_discrete_list, dim=1)

    def to_log_probs_tuple(self) -> LogProbsTuple:
        """
        Returns a LogProbsTuple. Only adds if tensor is not None. Otherwise,
        LogProbsTuple uses a default.
        """
        log_probs_tuple = LogProbsTuple()
        if self.continuous_tensor is not None:
            continuous = ModelUtils.to_numpy(self.continuous_tensor)
            log_probs_tuple.add_continuous(continuous)
        if self.discrete_list is not None:
            discrete = ModelUtils.to_numpy(self.discrete_tensor)
            log_probs_tuple.add_discrete(discrete)
        return log_probs_tuple

    def _to_tensor_list(self) -> List[torch.Tensor]:
        """
        Returns the tensors in the ActionLogProbs as a flat List of torch Tensors. This
        is private and serves as a utility for self.flatten()
        """
        tensor_list: List[torch.Tensor] = []
        if self.continuous_tensor is not None:
            tensor_list.append(self.continuous_tensor)
        if self.discrete_list is not None:
            tensor_list.append(self.discrete_tensor)
        return tensor_list

    def flatten(self) -> torch.Tensor:
        """
        A utility method that returns all log probs in ActionLogProbs as a flattened tensor.
        This is useful for algorithms like PPO which can treat all log probs in the same way.
        """
        return torch.cat(self._to_tensor_list(), dim=1)

    @staticmethod
    def from_buffer(buff: AgentBuffer) -> "ActionLogProbs":
        """
        A static method that accesses continuous and discrete log probs fields in an AgentBuffer
        and constructs the corresponding ActionLogProbs from the retrieved np arrays.
        """
        continuous: torch.Tensor = None
        discrete: List[torch.Tensor] = None  # type: ignore

        if BufferKey.CONTINUOUS_LOG_PROBS in buff:
            continuous = ModelUtils.list_to_tensor(buff[BufferKey.CONTINUOUS_LOG_PROBS])
        if BufferKey.DISCRETE_LOG_PROBS in buff:
            discrete_tensor = ModelUtils.list_to_tensor(
                buff[BufferKey.DISCRETE_LOG_PROBS]
            )
            # This will keep discrete_list = None which enables flatten()
            if discrete_tensor.shape[1] > 0:
                discrete = [
                    discrete_tensor[..., i] for i in range(discrete_tensor.shape[-1])
                ]
        return ActionLogProbs(continuous, discrete, None)


import numpy as np, json
import pickle, sys, argparse
import keras
from keras.models import Model
from keras import backend as K
from keras import initializers
from keras.optimizers import RMSprop
from keras.utils import to_categorical
from keras.callbacks import EarlyStopping, Callback, ModelCheckpoint
from keras.layers import *
from sklearn.metrics import classification_report, confusion_matrix, precision_recall_fscore_support, accuracy_score, f1_score
global seed
seed = 1337
np.random.seed(seed)
import gc
from sklearn.metrics import mean_squared_error,mean_absolute_error
from scipy.stats import pearsonr
from scipy.spatial.distance import cosine
#=============================================================
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
#=============================================================
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.log_device_placement = True
set_session(tf.Session(config=config))
#=============================================================

def attention(att_type, x, y):

    if att_type == 'simple':
        m_dash = dot([x, y], axes=[2,2])
        m = Activation('softmax')(m_dash)
        h_dash = dot([m, y], axes=[2,1])
        return multiply([h_dash, x])

    elif att_type == 'gated':
        alpha_dash = dot([y, x], axes=[2,2])
        alpha = Activation('softmax')(alpha_dash)
        x_hat = Permute((2, 1))(dot([x, alpha], axes=[1,2]))
        return multiply([y, x_hat])

    else:
        print ('Attention type must be either simple or gated.')

def emotionClass(testLabel):
    trueLabel     = []
    for i in range(testLabel.shape[0]):
        maxLen       = []
        for j in range(testLabel.shape[1]):
            temp = np.zeros((1,7),dtype=int)[0]
            pos  = np.nonzero(testLabel[i][j])[0]
            temp[pos] = 1
            maxLen.append(temp)
        trueLabel.append(maxLen)
    trueLabel = np.array(trueLabel)
    return trueLabel

def seventhClass(inputLabel, mask):
    updatedLabel = np.zeros((inputLabel.shape[0],inputLabel.shape[1],7), dtype ='float')
    for i in range(inputLabel.shape[0]):
        for j in range(list(mask[i]).count(1)):
            suM = np.sum(inputLabel[i][j])
            if suM == 0:
                updatedLabel[i][j][6] = 1
            else:
                updatedLabel[i][j][0:6] = inputLabel[i][j]
                updatedLabel[i,j,np.nonzero(updatedLabel[i][j])[0]]=1
    return updatedLabel

def featuresExtraction():
    global train_text, train_audio, train_video, senti_train_label, emo_train_label, train_mask
    global valid_text, valid_audio, valid_video, senti_valid_label, emo_valid_label, valid_mask
    global test_text, test_audio, test_video, senti_test_label, emo_test_label, test_mask
    global max_segment_len

    text          = np.load('MOSEI/text.npz',mmap_mode='r')
    audio         = np.load('MOSEI/audio.npz',mmap_mode='r')
    video         = np.load('MOSEI/video.npz',mmap_mode='r')

    train_text    = text['train_data']
    train_audio   = audio['train_data']
    train_video   = video['train_data']

    valid_text    = text['valid_data']
    valid_audio   = audio['valid_data']
    valid_video   = video['valid_data']

    test_text     = text['test_data']
    test_audio    = audio['test_data']
    test_video    = video['test_data']

    senti_train_label   = video['trainSentiLabel']
    senti_valid_label   = video['validSentiLabel']
    senti_test_label    = video['testSentiLabel']

    senti_train_label   = to_categorical(senti_train_label >= 0)
    senti_valid_label   = to_categorical(senti_valid_label >= 0)
    senti_test_label    = to_categorical(senti_test_label >= 0)

    emo_train_label   = video['trainEmoLabel']
    emo_valid_label   = video['validEmoLabel']
    emo_test_label    = video['testEmoLabel']

    train_length  = video['train_length']
    valid_length  = video['valid_length']
    test_length   = video['test_length']

    max_segment_len = train_text.shape[1]

    train_mask = np.zeros((train_video.shape[0], train_video.shape[1]), dtype='float')
    valid_mask = np.zeros((valid_video.shape[0], valid_video.shape[1]), dtype='float')
    test_mask  = np.zeros((test_video.shape[0], test_video.shape[1]), dtype='float')

    for i in xrange(len(train_length)):
        train_mask[i,:train_length[i]]=1.0

    for i in xrange(len(valid_length)):
        valid_mask[i,:valid_length[i]]=1.0

    for i in xrange(len(test_length)):
        test_mask[i,:test_length[i]]=1.0

    #====================== Add 7th class =========================================
    trainL = seventhClass(emo_train_label, train_mask)
    validL = seventhClass(emo_valid_label, valid_mask)
    testL  = seventhClass(emo_test_label, test_mask)

    #=================== Add multilabel class =====================================
    emo_train_label = emotionClass(trainL)
    emo_valid_label = emotionClass(validL)
    emo_test_label  = emotionClass(testL)

#=================================================================================
def calc_valid_result_emotion(result, test_label, test_mask):

    true_label=[]
    predicted_label=[]

    for i in range(result.shape[0]):
        for j in range(result.shape[1]):
            if test_mask[i,j]==1:
                true_label.append(test_label[i,j])
                predicted_label.append(result[i,j])
    true_label      = np.array(true_label)
    predicted_label = np.array(predicted_label)

    return true_label, predicted_label

def calc_valid_result_sentiment(result, test_label, test_mask):

    true_label=[]
    predicted_label=[]

    for i in range(result.shape[0]):
        for j in range(result.shape[1]):
            if test_mask[i,j]==1:
                true_label.append(np.argmax(test_label[i,j] ))
                predicted_label.append(np.argmax(result[i,j] ))
    return true_label, predicted_label

def weighted_accuracy(y_true, y_pred):
    TP, TN, FN, FP, N, P = 0, 0, 0, 0, 0, 0
    for i,j in zip(y_true,y_pred):
        if i == 1 and i == j:
            TP += 1
        elif i == 0 and i == j:
           TN += 1

        if i == 1 and i != j:
            FN += 1
        elif i == 0 and i != j:
           FP += 1

        if i == 1:
            P += 1
        else:
           N += 1

    w_acc = (1.0 * TP * (N / (1.0 * P)) + TN) / (2.0 * N)

    return w_acc, TP, TN, FP, FN, P, N
#=================================================================================
class MetricsCallback(keras.callbacks.Callback):
    def __init__(self, test_data):
        #super().__init__()
        self.test_data = test_data

    def on_train_begin(self, logs={}):
        self.Precision_senti    = []
        self.Recall_senti       = []
        self.Fscore_senti       = []
        self.Accuracy_senti     = []
        self.Weighted_acc_senti = []
        self.Precision_emo      = []
        self.Recall_emo         = []
        self.Fscore_emo         = []
        self.Accuracy_emo       = []
        self.Weighted_acc_emo   = []

    def on_epoch_end(self, epoch, logs={}):
        x_data    = self.test_data[0]
        y_actual  = self.test_data[1]

        #=============================== classification for sentiment ============================

        y_prediction = self.model.predict(x_data)

        true_label_senti, predicted_label_senti = calc_valid_result_sentiment(y_prediction[0], y_actual[0], test_mask)
        w_acc, TP, TN, FP, FN, P, N = weighted_accuracy(true_label_senti, predicted_label_senti)
        open('results/'+modality+'_senti.txt', 'a').write(str(epoch)+'\t'+
                                      str(attn_type)+'\t'+
                                      str(accuracy_score(true_label_senti, predicted_label_senti))+'\t' +
                                      str(precision_recall_fscore_support(true_label_senti, predicted_label_senti, average='weighted')[2])+'\t'+
                                      str(w_acc) + '\t('+
                                      str(TP) + ',' + str(TN) + ',' + str(FP) + ',' + str(FN) + ',' + str(P) + ',' + str(N) + ')\n')

        #=============================== classification for Emotion ============================
        th=[0.10,0.15,0.16,0.17,0.18,0.19,0.20,0.21,0.22,0.23,0.24,0.25,0.30,0.35,0.40,0.50]
        for t in range(len(th)):
            print th[t]
            emotion = ['Anger', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'No Class']
            y_prediction = self.model.predict(x_data)

            y_prediction[1][y_prediction[1] >= th[t]]  = 1
            y_prediction[1][y_prediction[1] <  th[t]]  = 0

            true_label_emo, predicted_label_emo = calc_valid_result_emotion(y_prediction[1], y_actual[1], test_mask)

            Acc_emo     = []
            F1Score_emo = []
            wAcc_emo    = []
            for i in range(7):
                Acc_emo.append(accuracy_score(true_label_emo[:,i], predicted_label_emo[:,i]))
                F1Score_emo.append(precision_recall_fscore_support(true_label_emo[:,i], predicted_label_emo[:,i], average='weighted')[2])
                wAcc_emo.append(weighted_accuracy(true_label_emo[:,i], predicted_label_emo[:,i])[0])

                w_acc, TP, TN, FP, FN, P, N = weighted_accuracy(true_label_emo[:,i], predicted_label_emo[:,i])
                open('results/'+modality+'_emo.txt', 'a').write('Threshold: ' + str(th[t]) +'\t'+
                                                  str(epoch)+'\t'+
                                                  str(attn_type)+'\t'+
                                                  str(accuracy_score(true_label_emo[:,i], predicted_label_emo[:,i]))+'\t' +
                                                  str(precision_recall_fscore_support(true_label_emo[:,i], predicted_label_emo[:,i], average='weighted')[2])+'\t'+
                                                  str(w_acc) + '\t('+
                                                  str(TP) + ',' + str(TN) + ',' + str(FP) + ',' + str(FN) + ',' + str(P) + ',' + str(N) + ')\t'+
                                                  str(emotion[i])+'\n')

            open('results/'+modality+'_emo.txt', 'a').write('Threshold: ' + str(th[t]) +'\t'+
                                              str(epoch)+'\t'+
                                              str(attn_type)+'\t'+
                                              str(np.average(Acc_emo)) + '\t'+
                                              str(np.average(F1Score_emo))+ '\t'+
                                              str(np.average(wAcc_emo)) + '\t()'+'\taverage\n')

def multimodal_cross_attention(attn_type, recurrent, timedistributed):

    featuresExtraction()
    global modality, emoName
    modality = 'trimodal'
    emoName  = ['Anger','Disgust','Fear','Happy','Sad','Surprise']

    runs = 1

    for run in range(runs):
        ################################### model architecture #################################
        in_text = Input(shape=(train_text.shape[1], train_text.shape[2]))
        in_audio = Input(shape=(train_audio.shape[1], train_audio.shape[2]))
        in_video = Input(shape=(train_video.shape[1], train_video.shape[2]))
        ########### masking layer ############
        masked_text = Masking(mask_value=0)(in_text)
        masked_audio = Masking(mask_value=0)(in_audio)
        masked_video = Masking(mask_value=0)(in_video)
        ########### recurrent layer ############
        drop0 = 0.3
        drop1 = 0.3
        r_drop = 0.3
        r_units = 200

        if recurrent:
            rnn_text = Bidirectional(GRU(r_units, return_sequences=True, dropout=r_drop, recurrent_dropout=r_drop),
                                     merge_mode='concat')(masked_text)
            rnn_audio = Bidirectional(GRU(r_units, return_sequences=True, dropout=r_drop, recurrent_dropout=r_drop),
                                      merge_mode='concat')(masked_audio)
            rnn_video = Bidirectional(GRU(r_units, return_sequences=True, dropout=r_drop, recurrent_dropout=r_drop),
                                      merge_mode='concat')(masked_video)

            inter_text = Dropout(drop0)(rnn_text)
            inter_audio = Dropout(drop0)(rnn_audio)
            inter_video = Dropout(drop0)(rnn_video)
        else:
            inter_text = Dropout(drop0)(masked_text)
            inter_audio = Dropout(drop0)(masked_audio)
            inter_video = Dropout(drop0)(masked_video)

        ########### timedistributed dense layer ############
        td_units = 100

        if timedistributed:
            td_text = Dropout(drop1)(TimeDistributed(Dense(td_units, activation='relu'))(inter_text))
            td_audio = Dropout(drop1)(TimeDistributed(Dense(td_units, activation='relu'))(inter_audio))
            td_video = Dropout(drop1)(TimeDistributed(Dense(td_units, activation='relu'))(inter_video))
        else:
            td_text = inter_text
            td_audio = inter_audio
            td_video = inter_video

        ########### attention layer ############
        ## cross modal cross utterance attention ##
        if attn_type == 'CIM':
            va_att = attention('simple', td_video, td_audio)
            vt_att = attention('simple', td_video, td_text)
            av_att = attention('simple', td_audio, td_video)
            at_att = attention('simple', td_audio, td_text)
            tv_att = attention('simple', td_text, td_video)
            ta_att = attention('simple', td_text, td_audio)

            merged = concatenate([va_att, vt_att, av_att, at_att, tv_att, ta_att, td_video, td_audio, td_text])

        ## uni modal cross utterance attention ##
        elif attn_type == 'ummu':
            vv_att = attention('simple', td_video, td_video)
            tt_att = attention('simple', td_text, td_text)
            aa_att = attention('simple', td_audio, td_audio)

            merged = concatenate([aa_att, vv_att, tt_att, td_video, td_audio, td_text])

        ## no attention ##
        elif attn_type == 'None':
            merged = concatenate([td_video, td_audio, td_text])
        else:
            print ("attn type must be either 'CIM' or 'ummu' or 'cmuu' or 'None'.")

        ########### output layer ############
        output_sentiment = TimeDistributed(Dense(2, activation='softmax'), name='output_sentiment')(merged)
        output_emotion   = TimeDistributed(Dense(7, activation='sigmoid'), name='output_emotion')(merged)
        model = Model([in_text, in_audio, in_video], [output_sentiment, output_emotion])
        model.compile(optimizer='adam', loss={'output_sentiment':'categorical_crossentropy', 'output_emotion':'binary_crossentropy'}, sample_weight_mode='temporal', metrics = {'output_sentiment': 'accuracy','output_emotion': 'accuracy'})

        ###################### model training #######################
        np.random.seed(run)
        path1 = 'weights/sentiment_'+modality+'_'+str(run)+'.hdf5'
        path2 = 'weights/emotion_'+modality+'_'+str(run)+'.hdf5'

        earlyStop_sentiment = EarlyStopping(monitor='val_output_sentiment_loss', patience=20)
        earlyStop_emotion = EarlyStopping(monitor='val_output_emotion_loss', patience=20)
        bestModel_sentiment = ModelCheckpoint(path1, monitor='val_output_sentiment_acc', verbose=1, save_best_only=True, mode='max')
        bestModel_emotion = ModelCheckpoint(path2, monitor='val_output_emotion_acc', verbose=1, save_best_only=True, mode='max')

        metrics_callback = MetricsCallback(test_data=([test_text, test_audio, test_video], [senti_test_label, emo_test_label]))
        history = model.fit([train_text, train_audio, train_video], [senti_train_label, emo_train_label],
                            epochs=50,
                            batch_size=16,
                            sample_weight=[train_mask, train_mask],
                            shuffle=True,
                            callbacks=[metrics_callback, bestModel_sentiment, bestModel_emotion],
                            validation_data=([valid_text, valid_audio, valid_video], [senti_valid_label,emo_valid_label], [valid_mask,valid_mask]),
                            verbose=1)

if __name__ == "__main__":
    global attn_type
    attn_type = 'CIM'
    multimodal_cross_attention(attn_type=attn_type, recurrent=True, timedistributed=True)


import tensorflow as tf
import tensorflow.contrib.slim as slim
import numpy as np
import sys
import random
# from contact_point_dataset_torch_multi_label import MyDataset
from simple_dataset import MyDataset
import os
import time
import argparse
from functools import partial
import time

from torch.utils.data import DataLoader
import torch
# from torch.utils.tensorboard import SummaryWriter
from tensorboardX import SummaryWriter
import datetime

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UTILS_DIR = os.path.abspath(os.path.join(BASE_DIR, '..', 'utils'))
sys.path.append(UTILS_DIR)
from data_helper import *
from coord_helper import *
from train_helper import *
from bullet_helper import *
from rotation_lib import *
import s1_model_multi_label as s1_model
import s3_classifier_model as s3_model
from s2_utils import *

import s3_replay_buffer_pose as ReplayBuffer
import ES_multithread
import multiprocessing
from scipy.special import softmax

def restore_model_s1(epoch, save_top_dir, sess):
	ckpt_path = os.path.join(save_top_dir,str(epoch)+'model.ckpt')
	variables = slim.get_variables_to_restore()
	variables_to_restore = [v for v in variables if 's2a_' not in v.name and ('s1_' in v.name)]
	# for v in variables_to_restore:
		# print(v.name)
	saver = tf.train.Saver(variables_to_restore)
	print("restoring from %s" % ckpt_path)
	saver.restore(sess, ckpt_path)

def calc_pose_cem_init(args, pc_o_all, pc_h_all, transl_s1_all, aa_s1_all):
	b_size = transl_s1_all.shape[0]
	cem_init_transl = np.zeros((b_size, 3))
	cem_init_aa = np.zeros((b_size, 3))
	cem_rotation_center_o = np.zeros((b_size, 3))
	#pc_tmp_o_all = np.zeros((b_size, args.s3_num_cp, 3))
	#pc_tmp_h_all = np.zeros((b_size, args.s3_num_cp, 3))
	#pc_tmp_rotated_s1_o_all = np.zeros((b_size, args.s3_num_cp, 3))

	for bi in range(b_size):
		transl_s1 = transl_s1_all[bi]
		aa_s1 = aa_s1_all[bi]

		cem_init_transl[bi] = transl_s1
		cem_init_aa[bi] = aa_s1

		#cem_rotation_center_o[bi] = np.mean(pc_o_all[bi], axis=0)
		# debug info
		#pc_tmp_o_all[bi] = pc_tmp_o
		#pc_tmp_rotated_s1_o_all[bi] = pc_tmp_rotated_s1_o
		#pc_tmp_h_all[bi] = pc_tmp_h

	info_dict = {
		#'s3_partial_pc_o': np.copy(pc_tmp_o_all),
		#'s3_partial_pc_rotated_s1_o': np.copy(pc_tmp_rotated_s1_o_all),
		#'s3_partial_pc_h': np.copy(pc_tmp_h_all),
		'cem_rotation_center_o': np.copy(cem_rotation_center_o),
	}
	return cem_init_transl, cem_init_aa, cem_rotation_center_o, info_dict

from s3_rl_collect import cem_transform_pc_batch, cem_eval, bullet_check
import s3_bullet_checker as bullet_checker
import s3_bullet_checker_eval as bullet_checker_eval

def train(args, train_set, train_loader, test_set, test_loader, writer, result_folder, file_name, extra_dict=None):
	model_folder = os.path.join(result_folder, 'models')
	can_write = not (writer is None)

	# stage 1
	pc_o_pl, pc_h_pl, z_pl, gt_transl_pl, gt_aa_pl, pose_mult_pl = s1_model.placeholder_inputs(args.batch_size, 4096, args)
	pred_transl_tf, pred_aa_tf, end_points_s1 = s1_model.get_model(pc_o_pl, pc_h_pl, z_pl)
	loss_transl_tf, loss_aa_tf, min_pose_idx_tf = s1_model.get_loss(pred_transl_tf, pred_aa_tf, gt_transl_pl, gt_aa_pl, pose_mult_pl, float(args.loss_transl_const), end_points_s1)
	loss_s1_tf = float(args.loss_transl_const) * loss_transl_tf + loss_aa_tf

	#stage 3
	pc_combined_pl, gt_succ_label_pl = s3_model.placeholder_inputs(args.batch_size, 4096, with_normal=False, args=args)
	pred_succ_cla_score_tf, end_points = s3_model.get_model(pc_combined_pl)

	pred_succ_cla_tf = tf.math.argmax(pred_succ_cla_score_tf, axis=-1)
	# loss_succ_cla_tf = s3_model.get_loss(pred_succ_cla_score_tf, gt_succ_label_pl, end_points)

	# loss_tf = loss_listnet_tf + loss_ce_tf
	# loss_tf = tf.boolean_mask(loss_tf, non_nan_mask_pl)
	# loss_tf = tf.reduce_mean(loss_tf)
	# print('loss tf', loss_tf)
	# train_op = tf.train.AdamOptimizer(learning_rate=args.learning_rate).minimize(loss_tf)

	init_op = tf.group(tf.global_variables_initializer(),
				tf.local_variables_initializer())

	config = tf.ConfigProto()
	config.gpu_options.allow_growth = True

	sess = tf.Session(config=config)
	sess.run(init_op)
	saver = tf.train.Saver(max_to_keep=1000)
	loss_tracker = LossTracker()
	loss_tracker_test = LossTracker()

	epoch_init = 0
	if not train_loader is None:
		epoch_iter = len(train_loader)

	if args.pretrain_s1:
		pretrain_s1_folder = os.path.abspath(os.path.join(result_folder, '..', '..', args.pretrain_s1_folder, args.pretrain_s1_model_name, 'models'))
		restore_model_s1(args.pretrain_s1_epoch, pretrain_s1_folder, sess)

	if args.pretrain_s3:
		pretrain_s3_folder = os.path.abspath(os.path.join(result_folder, '..', '..', args.pretrain_s3_folder, args.pretrain_s3_model_name, 'models'))

		if args.pretrain_s3_folder_dir != '':
			pretrain_s3_folder = args.pretrain_s3_folder_dir
		restore_model_s3(args.pretrain_s3_epoch, pretrain_s3_folder, sess)
	else:
		print(args.s3_model_dir)
		tmp = restore_model_s3_second_last(args.s3_model_dir, sess)
		assert tmp # make sure that one model is restored

	if args.restore_model_epoch != -1:
		epoch_init = args.restore_model_epoch
		restore_model_folder = os.path.abspath(os.path.join(result_folder, '..', args.restore_model_name, 'models'))
		restore_model_generic(epoch_init, restore_model_folder, saver, sess)

	total_ct = 0

	cem = ES_multithread.Searcher(
        action_dim=6,
        max_action=float(args.cem_max_transl),
        max_action_aa=float(args.cem_max_aa),
        sigma_init=float(args.cem_sigma_init_transl),
        sigma_init_aa=float(args.cem_sigma_init_aa),
        pop_size=args.cem_pop_size,
        damp=float(args.cem_damp_transl),
        damp_limit=float(args.cem_damp_limit_transl),
        damp_aa=float(args.cem_damp_aa),
        damp_limit_aa=float(args.cem_damp_limit_aa),
        parents=args.cem_parents,
	)

	fcl_hook_dict = extra_dict['fcl_hook_dict']
	fcl_object_dict = extra_dict['fcl_object_dict']

	if not args.no_bullet_check:
		pool = multiprocessing.Pool(processes=args.batch_size)
		p_list = pool.map(partial(p_init_multithread, gui=args.bullet_gui), range(args.batch_size))

	if not args.run_test:
		saved_buffer_name_list = [get_2nd_last_dir(args.s3_buffer_dir)]
		replay_buffer = ReplayBuffer.combine_replay_buffer(saved_buffer_name_list, train_set, args.home_dir_data, preload_data=False)

	eval_folder_dir = os.path.join(result_folder, 'eval')
	mkdir_if_not(eval_folder_dir)
	# for epoch_i in range(args.max_epochs):
	epoch_i = 0
	tested = False
	while True:
		if epoch_i == args.max_epochs:
			break
		loss_tracker.reset()

		run_test = False
		if (not args.no_eval) and (((epoch_i + 1) % args.eval_epoch_freq == 0) or args.run_test):
			if not tested:
				run_test = True
				epoch_i -= 1
				tested = True
			if tested:
				tested = False
		if args.run_test:
			epoch_i = 0
		if not run_test:
			loader = train_loader
			dataset = train_set
		else:
			loader = test_loader
			dataset = test_set

		info_dict_all = {}

		for i, batch_dict in enumerate(loader):
			total_ct += 1

			if not args.no_bullet_check:
				if (total_ct + 1) % 20 == 0:
					print('reset')
					p_list = pool.map(partial(p_reset_multithread, p_list=p_list, gui=args.bullet_gui), range(args.batch_size))

			log_it = ((total_ct % args.log_freq ) == 0) and can_write
			pc_o = batch_dict['input1']
			pc_h = batch_dict['input2']
			b_size = pc_o.shape[0]
			object_urdf = batch_dict['urdf_o']
			hook_urdf = batch_dict['urdf_h']
			result_file_name = batch_dict['result_file_name']

			fcl_hook_model = [fcl_hook_dict[name] for name in batch_dict['hook_name']]
			fcl_object_model = [fcl_object_dict[name] for name in batch_dict['object_name']]

			if run_test:
				eval_sample_n = args.eval_sample_n
			else:
				eval_sample_n = 1
			info_dict = {}
			bullet_succ = np.zeros((b_size, eval_sample_n))
			for i_eval in range(args.eval_sample_n):
				z = np.random.normal(size=(pc_o.shape[0],1,32))
				feed_dict_s1 = {
					pc_o_pl: pc_o[:, :, :3],
					pc_h_pl: pc_h[:, :, :3],
					z_pl: z,
				}

				pred_transl, pred_aa = sess.run([
					pred_transl_tf, pred_aa_tf
				], feed_dict=feed_dict_s1)

				# stage 3
				cem_init_transl, cem_init_aa, cem_rotation_center_o, cem_info_dict = calc_pose_cem_init(
					args=args,
					pc_o_all=pc_o[:, :, :3],
					pc_h_all=pc_h[:, :, :3],
					transl_s1_all=pred_transl,
					aa_s1_all=pred_aa,
				)

				pc_o_cem_init = transform_pc_batch(pc_o[:, :, :3], cem_init_transl, cem_init_aa)

				cem_eval_partial = partial(cem_eval,
					pc_o=pc_o_cem_init,
					pc_h=pc_h[:, :, :3],
					rotation_center_o=cem_rotation_center_o,
					sess_tf=sess,
					pc_combined_pl=pc_combined_pl,
					pred_succ_cla_score_tf=pred_succ_cla_score_tf,
				)

				cem_search_info_dict = {
					'cem_elite_pose': np.zeros((b_size, args.cem_n_iter, 6)),
					'cem_elite_pose_scores': np.zeros((b_size, args.cem_n_iter))
				}
				cem_max_score = np.zeros((b_size))
				cem_out_transl = np.zeros((b_size, 3))
				cem_out_aa = np.zeros((b_size, 3))
				cem_out_pose = np.zeros((b_size, 6))
				for ii in range(args.cem_run_n):
					_, cem_out_pose_tmp, cem_score_tmp, cem_search_info_dict_tmp = cem.search(
						b_size,
						np.array([[0, 0, 0, 1e-6, 0, 0]] * b_size),
						cem_eval_partial,
						n_iter=args.cem_n_iter,
						elitism=True,
						visualize=False,
						visualize_func=None,
					)
					for jj in range(b_size):
						cur_score = cem_score_tmp[jj]
						if cem_max_score[jj] < cur_score:
							cem_max_score[jj] = cur_score
							cem_out_transl[jj] = cem_out_pose_tmp[jj, :3]
							cem_out_aa[jj] = cem_out_pose_tmp[jj, 3:]
							cem_out_pose[jj] = cem_out_pose_tmp[jj]
							cem_search_info_dict['cem_elite_pose'][jj] = cem_search_info_dict_tmp['cem_elite_pose'][jj]
							cem_search_info_dict['cem_elite_pose_scores'][jj] = cem_search_info_dict_tmp['cem_elite_pose_scores'][jj]
				# convert the cem output pose to object pose
				pc_o_cem = cem_transform_pc_batch(pc_o_cem_init[:, :, :3], cem_rotation_center_o, cem_out_transl, cem_out_aa)

				final_pred_transl, final_pred_aa = best_fit_transform_batch(pc_o[:, :, :3], pc_o_cem)

				if not args.no_bullet_check:
					bullet_check_one_pose = bullet_checker_eval.check_one_pose_simple if args.use_bullet_checker else bullet_checker.check_one_pose_simple
					# bullet check
					bullet_check_func = partial(
						bullet_check,
						bullet_check_one_pose=bullet_check_one_pose,
						transl=final_pred_transl,
						aa=final_pred_aa,
						p_list=p_list,
						result_file_name=result_file_name,
						hook_urdf=hook_urdf,
						object_urdf=object_urdf,
						fcl_hook_model=fcl_hook_model,
						fcl_object_model=fcl_object_model,
						gui=args.bullet_gui,
					)

					for bi, (flag_tmp, bullet_final_transl, bullet_final_quat) in enumerate(pool.imap(bullet_check_func, range(b_size))):
						succ = 1. if flag_tmp else 0.
						bullet_succ[bi, i_eval] = succ

						hook_name, object_name = split_result_file_name(result_file_name[bi])
						if not run_test:
							replay_buffer.add(object_name, hook_name, np.append(final_pred_transl[bi], final_pred_aa[bi]), succ)

				for ii in range(b_size):
					if i_eval == 0:
						info_dict[result_file_name[ii]] = []

					info_dict[result_file_name[ii]].append({
						'z': z[ii].tolist(),
						's1_transl': pred_transl[ii].tolist(),
						's1_aa': pred_aa[ii].tolist(),
						'cem_init_transl': cem_init_transl[ii].tolist(),
						'cem_init_aa': cem_init_aa[ii].tolist(),

						'succ': bullet_succ[ii, i_eval],
						'final_pred_transl': final_pred_transl[ii].tolist(),
						'final_pred_aa': final_pred_aa[ii].tolist(),
						'cem_out_transl': cem_out_pose[ii, :3].tolist(),
						'cem_out_aa': cem_out_pose[ii, 3:].tolist(),
						# 'cem_elite_pose': cem_elite_pose[ii].tolist()
					})
					for tmp_key in cem_info_dict:
						info_dict[result_file_name[ii]][-1][tmp_key] = cem_info_dict[tmp_key][ii].tolist()

					for tmp_key in cem_search_info_dict:
						info_dict[result_file_name[ii]][-1][tmp_key] = cem_search_info_dict[tmp_key][ii].tolist()

			info_dict_all.update(info_dict)

			loss_dict = {
				'bullet_succ_acc': np.mean(bullet_succ),
				'bullet_succ_acc_max': np.mean(np.max(bullet_succ, axis=-1))
			}
			loss_tracker.add_dict(loss_dict)

			if log_it:
				write_tb(loss_dict, writer, 'test' if run_test else 'train', total_ct)
			print('epoch {} iter {}/{} {}'.format(epoch_i, i, epoch_iter, loss_dict_to_str(loss_dict)))
			loss_dict_epoch = loss_tracker.stat()
			print('cumulative, epoch {} {} {}'.format(epoch_i, 'test' if run_test else 'train', loss_dict_to_str(loss_dict_epoch)))

			# periodically save buffer
			if int(total_ct) % args.s3_buffer_freq == 0 and (not run_test):
				print('save buffer', args.s3_buffer_dir, replay_buffer.buffer_succ.size, replay_buffer.buffer_fail.size, )
				replay_buffer.save_pkl(args.s3_buffer_dir, replay_buffer)

				# save info dict
				info_dict_dir = os.path.join(eval_folder_dir, '{}_eval_epoch_{}_ct_{}_{}.json'.format(file_name, str(epoch_i + 1), int(total_ct), 'test' if run_test else 'train'))
				save_json(info_dict_dir, info_dict_all)

			# periodically load s3 model
			if int(total_ct) % args.s3_model_freq == 0 and (not run_test):
				restore_model_s3_second_last(args.s3_model_dir, sess)

			if (total_ct % args.model_save_freq == 0) and not args.no_save and (not run_test):
				save_model_generic(epoch_init + total_ct, model_folder, saver, sess)

			if total_ct % 5 == 0 or total_ct == epoch_iter:
				# save info dict
				info_dict_dir = os.path.join(eval_folder_dir, '{}_eval_epoch_{}_{}.json'.format(file_name, str(epoch_i + 1), 'test' if run_test else 'train'))
				save_json(info_dict_dir, info_dict_all)

		loss_dict_epoch = loss_tracker.stat()
		if can_write:
			write_tb(loss_dict_epoch, writer, 'test_epoch' if run_test else 'train_epoch', total_ct)
			print('epoch {} {} {}'.format(epoch_i, 'test' if run_test else 'train', loss_dict_to_str(loss_dict_epoch)))

		epoch_i += 1

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument("--home_dir_data", default="../data")
	parser.add_argument('--pointset_dir', default='/scr2/')
	parser.add_argument('--bohg4', action='store_true')
	parser.add_argument('--no_vis', action='store_true')

	parser.add_argument('--model_name', default='s3_rl_collect')
	parser.add_argument('--comment', default='')
	parser.add_argument('--exp_name', default='exp_s3')
	parser.add_argument('--debug', action='store_true')
	parser.add_argument('--log_freq', type=int, default=2)

	parser.add_argument('--train_list', default='train_list')
	parser.add_argument('--test_list', default='test_list')
	parser.add_argument('--restrict_object_cat', default='')

	parser.add_argument('--run_test', action='store_true')
	parser.add_argument('--no_save', action='store_true')
	parser.add_argument('--overfit', action='store_true')
	parser.add_argument('--restore_model_name', default='')
	parser.add_argument('--restore_model_epoch', type=int, default=-1)
	parser.add_argument('--max_epochs', type=int, default=10000)
	parser.add_argument('--eval_epoch_freq', type=int, default=2)
	parser.add_argument('--eval_sample_n', type=int, default=1)
	parser.add_argument('--model_save_freq', type=int, default=3000)
	parser.add_argument('--no_eval', action='store_true')

	parser.add_argument('--loss_transl_const', default=1)

	parser.add_argument('--data_one_pose', action='store_true')
	parser.add_argument('--data_vary_scale', action='store_true')
	parser.add_argument('--data_more_pose', action='store_true')
	parser.add_argument('--data_vary_scale_more_pose', action='store_true')

	parser.add_argument('--batch_size', type=int, default=10)
	parser.add_argument('--learning_rate', type=float, default=1e-4)

	#s1 argument
	parser.add_argument('--z_dim', type=int, default=32)

	#s2 argument
	parser.add_argument('--top_k_o', type=int, default=128)
	parser.add_argument('--top_k_h', type=int, default=128)
	parser.add_argument('--n_gt_sample', type=int, default=128)
	parser.add_argument('--top_k_corr', type=int, default=256)
	parser.add_argument('--pose_loss_l2', action='store_true')

	#s3 argument
	parser.add_argument('--s3_num_cp', type=int, default=3)
	parser.add_argument('--cem_run_n', type=int, default=1)
	parser.add_argument('--cem_n_iter', type=int, default=10)
	parser.add_argument('--cem_max_transl', default=0.02)
	parser.add_argument('--cem_max_aa', default=0.5)
	parser.add_argument('--cem_sigma_init_transl', default=1e-2)
	parser.add_argument('--cem_sigma_init_aa', default=1e-1)
	parser.add_argument('--cem_pop_size', type=int, default=32)
	parser.add_argument('--cem_damp_transl', default=0.005)
	parser.add_argument('--cem_damp_limit_transl', default=1e-2)
	parser.add_argument('--cem_damp_aa', default=0.1)
	parser.add_argument('--cem_damp_limit_aa', default=0.1)
	parser.add_argument('--cem_parents', type=int, default=10)

	parser.add_argument('--bullet_gui', action='store_true')
	parser.add_argument('--s3_train_folder_dir', default='/juno/downloads/new_hang_training/')
	parser.add_argument('--s3_train_name', default='s3')
	parser.add_argument('--s3_device_name', default='bohg4')
	parser.add_argument('--s3_buffer_dir', default='')
	parser.add_argument('--s3_model_dir', default='')
	parser.add_argument('--no_fcl', action='store_true')
	parser.add_argument('--s3_buffer_freq', default=1000, type=int)
	parser.add_argument('--s3_model_freq', default=1000, type=int)
	parser.add_argument('--use_bullet_checker', action='store_true')

	parser.add_argument('--pretrain_s1', action='store_true')
	parser.add_argument('--pretrain_s1_folder', default='exp_s1')
	parser.add_argument('--pretrain_s1_model_name', default='Feb17_23-51-13_s1_matching_model_loss_transl_const_1000')
	parser.add_argument('--pretrain_s1_epoch', default=40500, type=int)

	parser.add_argument('--pretrain_s3', action='store_true')
	parser.add_argument('--pretrain_s3_folder', default='exp_s3')
	parser.add_argument('--pretrain_s3_folder_dir', default='')
	parser.add_argument('--pretrain_s3_model_name', default='Feb19_14-26-47_s3_classifier_model_new_data')
	parser.add_argument('--pretrain_s3_epoch', default=750000, type=int)

	parser.add_argument('--lin', action='store_true')
	parser.add_argument('--no_bullet_check', action='store_true')
	parser.add_argument('--seed',default=3323,type=int)

	parser.add_argument('--parallel_n', default=-1, type=int)
	parser.add_argument('--parallel_id', default=-1, type=int)

	args = parser.parse_args()

	args.data_more_pose = True

	print("args.seed",args.seed)
	random.seed(args.seed)
	torch.manual_seed(args.seed)
	np.random.seed(args.seed)

	if args.no_bullet_check:
		args.no_fcl = True

	if args.bohg4:
		args.pointset_dir = '/scr1/yifan'
		args.no_vis = True
		args.home_dir_data = '/scr1/yifan/hang'

	if args.lin:
		args.s3_train_folder_dir = os.path.abspath('./new_hang_training')
		print(args.s3_train_folder_dir)

	file_name = "{}".format(args.model_name)
	file_name += '_{}'.format(args.restrict_object_cat) if args.restrict_object_cat != '' else ''
	file_name += "_{}".format(args.comment) if args.comment != "" else ""
	file_name += '_overfit' if args.overfit else ''
	folder_name = datetime.datetime.now().strftime('%b%d_%H-%M-%S_') + file_name

	if args.use_bullet_checker:
		args.no_fcl = True

	if args.run_test:
		folder_name += '_test'

	result_folder = 'runs/{}'.format(folder_name)
	if args.exp_name is not "":
		result_folder = 'runs/{}/{}'.format(args.exp_name, folder_name)
	if args.debug:
		result_folder = 'runs/debug/{}'.format(folder_name)

	model_folder = os.path.join(result_folder, 'models')

	if not os.path.exists(model_folder):
		os.makedirs(result_folder)

	if args.debug:
		args.s3_train_name = 'debug'
	else:
		args.s3_train_name += "_{}".format(args.comment) if args.comment != "" else ""
	if not args.run_test:
		s3_train_dir = os.path.join(args.s3_train_folder_dir, args.s3_train_name)
		mkdir_if_not(s3_train_dir)

		args.s3_buffer_dir = os.path.join(s3_train_dir, 'buffers')
		mkdir_if_not(args.s3_buffer_dir)
		args.s3_buffer_dir = os.path.join(args.s3_buffer_dir, args.s3_device_name)
		mkdir_if_not(args.s3_buffer_dir)

		args.s3_model_dir = os.path.join(s3_train_dir, 'models')
		mkdir_if_not(args.s3_model_dir)

	print("---------------------------------------")
	print("Model Name: {}, Train List: {}, Test List: {}".format(args.model_name, args.train_list, args.test_list))
	print("---------------------------------------")

	if args.run_test:
		print("Restore Model: {}, Test Sample N: {}".format(args.restore_model_name, args.eval_sample_n))
		print("---------------------------------------")

	writer = None
	if not args.run_test:
		writer = SummaryWriter(log_dir=result_folder, comment=file_name)

	#record all parameters value
	with open("{}/parameters.txt".format(result_folder), 'w') as file:
		for key in sorted(vars(args).keys()):
			value = vars(args)[key]
			file.write("{} = {}\n".format(key, value))

	cp_result_folder_dir = os.path.join(args.home_dir_data, 'dataset_cp')
	train_list_dir = os.path.join(cp_result_folder_dir, 'labels', '{}.txt'.format(args.train_list))
	test_list_dir = os.path.join(cp_result_folder_dir, 'labels', '{}.txt'.format(args.test_list))

	print('TRAIN_LIST:', args.train_list, train_list_dir)
	print('TEST_LIST:', args.test_list, test_list_dir)

	# if args.overfit:
		# args.no_eval = True
		# args.no_save = True
	if args.run_test:
		args.max_epochs = 1

	if args.restore_model_name != '':
		assert args.restore_model_epoch != -1

	if args.restore_model_epoch != -1:
		assert args.restore_model_name != ''

	train_loader = None
	if (not args.run_test) or args.overfit:
		train_set = MyDataset(args.home_dir_data, train_list_dir, is_train=True, use_partial_pc=True, use_fcl=(not args.no_fcl), split_n=args.parallel_n, split_id=args.parallel_id, args=args)
		train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,
							num_workers=6, collate_fn=MyDataset.pad_collate_fn_for_dict, drop_last=True)

	if not args.no_eval:
		test_set = MyDataset(args.home_dir_data, test_list_dir, is_train=False, use_partial_pc=True, use_fcl=(not args.no_fcl), split_n=args.parallel_n, split_id=args.parallel_id, args=args)

		test_loader = DataLoader(train_set if args.overfit else test_set, batch_size=args.batch_size, shuffle=True,
									num_workers=6, collate_fn=MyDataset.pad_collate_fn_for_dict, drop_last=True)
	else:
		test_set = None
		test_loader = None
	# if not args.run_test:
	# 	print('len of train {} len of test {}'.format(len(train_set), len(test_set)))
	# else:
	# 	print('len of train {} len of test {}'.format(len(test_set), len(test_set)))
	if args.run_test:
		train_set = test_set
		train_loader = test_loader
	extra_dict = {
		'fcl_object_dict': train_set.fcl_object_dict,
		'fcl_hook_dict': train_set.fcl_hook_dict,
	}
	if not args.run_test:
		train(args, train_set, train_loader, test_set, test_loader, writer, result_folder, file_name, extra_dict=extra_dict)
	else:
		train(args, test_set, test_loader, test_set, test_loader, writer, result_folder, file_name, extra_dict=extra_dict)


#!/usr/bin/env python3
# This file is covered by the LICENSE file in the root of this project.

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import imp
import yaml
import time
import __init__ as booger
import cv2
import os
import numpy as np

from backbones.config import *
from tasks.classification.modules.head import *
from tasks.classification.modules.classifier import *

class UserPytorch():
  def __init__(self, path):
    # parameters
    self.path = path

    # config from path
    try:
      yaml_path = self.path + "/cfg.yaml"
      print("Opening config file %s" % yaml_path)
      self.CFG = yaml.load(open(yaml_path, 'r'))
    except Exception as e:
      print(e)
      print("Error opening cfg.yaml file from trained model.")
      quit()

    # get the data
    parserModule = imp.load_source("parserModule",
                                   booger.TRAIN_PATH + '/tasks/classification/dataset/' +
                                   self.CFG["dataset"]["name"] + '/parser.py')
    self.parser = parserModule.Parser(img_prop=self.CFG["dataset"]["img_prop"],
                                      img_means=self.CFG["dataset"]["img_means"],
                                      img_stds=self.CFG["dataset"]["img_stds"],
                                      classes=self.CFG["dataset"]["labels"],
                                      train=False)

    # some useful data
    self.data_h, self.data_w, self.data_d = self.parser.get_img_size()
    self.means, self.stds = self.parser.get_means_stds()
    self.means = torch.tensor(self.means)
    self.stds = torch.tensor(self.stds)
    self.nclasses = self.parser.get_n_classes()

    # architecture definition
    # get weights?
    try:
      self.pytorch_path = os.path.join(self.path, "model.pytorch")
      self.model = torch.jit.load(self.pytorch_path)
      print("Successfully Pytorch-traced model from ", self.pytorch_path)
    except Exception as e:
      print("Couldn't load Pytorch-traced network. Error: ", e)
      quit()

    # GPU?
    self.gpu = False
    self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
      self.gpu = True
      cudnn.benchmark = True
      cudnn.fastest = True
      self.model.cuda()
      self.means = self.means.cuda()
      self.stds = self.stds.cuda()

  def infer(self, bgr_img, topk=1, verbose=True):
    # get sizes
    original_h, original_w, original_d = bgr_img.shape

    # resize
    bgr_img = cv2.resize(bgr_img, (self.data_w, self.data_h),
                         interpolation=cv2.INTER_LINEAR)

    # check if network is RGB or mono
    if self.CFG["dataset"]["img_prop"]["depth"] == 3:
      # get make rgb
      if verbose:
        print("Converting bgr to rgb")
      rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
    elif self.CFG["dataset"]["img_prop"]["depth"] == 1:
      # get grayscale
      if verbose:
        print("Converting to grayscale")
      rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2GRAY)
    else:
      raise NotImplementedError(
          "Network has to have 1 or 3 channels. Anything else must be implemented.")

    # to tensor
    rgb_tensor = torch.from_numpy(rgb_img)

    # to gpu
    if self.gpu:
      rgb_tensor = rgb_tensor.cuda()

    # permute and normalize
    rgb_tensor = (rgb_tensor.float() / 255.0 - self.means) / self.stds
    rgb_tensor = rgb_tensor.permute(2, 0, 1)

    # add batch dimension
    rgb_tensor = rgb_tensor.unsqueeze(0)

    # infer
    start = time.time()
    logits = self.model(rgb_tensor)
    if torch.cuda.is_available():
      torch.cuda.synchronize()
    max_k = logits[0].topk(topk, dim=-1, largest=True, sorted=True)
    time_to_infer = time.time() - start

    # print result, and put in output lists
    classes = []
    classes_str = []
    if verbose:
      print("Time to infer: {:.3f}s".format(time_to_infer))
    for i, idx in enumerate(max_k[1], 1):
      idx_cpu = idx.cpu().item()
      class_string = self.parser.get_class_string(idx_cpu)
      classes.append(idx_cpu)
      classes_str.append(class_string)
      if verbose:
        print("[{}]: {}, {}".format(i, idx_cpu, class_string))

    return classes, classes_str


#!/usr/bin/env python
#
# This file is part of the Emotions project. The complete source code is
# available at https://github.com/luigivieira/emotions.
#
# Copyright (c) 2016-2017, Luiz Carlos Vieira (http://www.luiz.vieira.nom.br)
#
# MIT License
#
# 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 os
from collections import OrderedDict
import numpy as np

from gabor import GaborBank
from data import FaceData
from faces import FaceDetector

from sklearn import svm
from sklearn.externals import joblib

#=============================================
class InvalidModelException(Exception):
    '''
    Exception indicating that the detection model could not be loaded (or didn't
    exist).
    '''
    pass

#=============================================
class EmotionsDetector:
    """
    Implements the detector of prototypic emotions on face images.
    """

    #---------------------------------------------
    def __init__(self):
        """
        Class constructor.
        """

        self._clf = svm.SVC(kernel='rbf', gamma=0.001, C=10,
                            decision_function_shape='ovr',
                            probability=True, class_weight='balanced')
        """
        Support Vector Machine with used as the model for the detection of the
        prototypic emotions. For details on the selection of the kernel and
        parameters, refer to the PhD thesis of the author of this code.
        """

        self._emotions = OrderedDict([
                             (0, 'neutral'), (1, 'happiness'), (2, 'sadness'),
                             (3, 'anger'), (4, 'fear'),  (5, 'surprise'),
                             (6, 'disgust')
                         ])
        """
        Class and labels of the prototypic emotions detected by this model.
        """

        modulePath = os.path.dirname(__file__)
        self._modelFile = os.path.abspath('{}/models/emotions_model.dat' \
                            .format(modulePath))
        """
        Name of the file used to persist the model in the disk.
        """

        # Load the model from the disk, if its file exists
        if not os.path.isfile(self._modelFile):
            raise InvalidModelException('Could not find model file: {}' \
                    .format(self._modelFile))

        if not self.load():
            raise InvalidModelException('Could not load model from file: {}' \
                    .format(self._modelFile))

    #---------------------------------------------
    def load(self):
        """
        Loads the SVM model from the disk.

        Returns
        -------
        ret: bool
            Indication on if the loading was succeeded or not.
        """

        try:
            clf = joblib.load(self._modelFile)
        except:
            return False

        self._clf = clf
        return True

    #---------------------------------------------
    def _relevantFeatures(self, gaborResponses, facialLandmarks):
        """
        Get the features that are relevant for the detection of emotions
        from the matrix of responses to the bank of Gabor kernels.

        The feature vector returned by this method can be used for training and
        predicting, using a linear SVM.

        Parameters
        ----------
        gaborResponses: numpy.array
            Matrix of responses to the bank of Gabor kernels applied to the face
            region of an image. The first dimension of this matrix has size 32,
            one for each kernel in the bank. The other two dimensions are in the
            same size as the original image used for their extraction.

        facialLandmarks: numpy.array
            Bidimensional matrix with the coordinates of each facial landmark
            detected in the face image from where the responses were obtained.

        Returns
        -------
        featureVector: list
            A list with the responses of the 32 kernels at each of the
            face landmarks.
        """

        # Get the 32 responses at the positions of all the face landmarks
        points = np.array(facialLandmarks)

        # Try to get the responses for all points. If an exception is caught,
        # it is because some landmarks are out of the image area (i.e. the face
        # is partially occluded, but it was still possible to detect). In this
        # case, assume 0.0 for the responses of the landmarks outside the image
        # area.
        try:
            responses = gaborResponses[:, points[:, 1], points[:, 0]]
        except:
            w = gaborResponses.shape[2]
            h = gaborResponses.shape[1]

            responses = np.zeros((32, 68), dtype=float)
            for i in range(len(points)):
                x = points[i][0]
                y = points[i][1]
                if x < w and y < h:
                    responses[:, i] = gaborResponses[:, y, x]
                else:
                    responses[:, i] = 0.0

        # Reshape the bi-dimensional matrix to a single dimension
        featureVector = responses.reshape(-1).tolist()

        return featureVector

    #---------------------------------------------
    def detect(self, face, gaborResponses):
        """
        Detects the emotions based on the given features.

        Parameters
        ----------
        face: FaceData
            Instance of the FaceData object with the facial landmarks detected
            on the facial image.
        gaborResponses: numpy.array
            Matrix of responses to the bank of Gabor kernels applied to the face
            region of an image. The first dimension of this matrix has size 32,
            one for each kernel in the bank. The other two dimensions are in the
            same size as the original image used for their extraction.

        Returns
        -------
        probabilities: OrderedDict
            The probabilities of each of the prototypic emotion, in format:
            {'anger': value, 'contempt': value, [...]}
        """

        # Filter only the responses at the facial landmarks
        features = self._relevantFeatures(gaborResponses, face.landmarks)

        # Return the prediction based on the given features
        return self.predict(features)

    #---------------------------------------------
    def predict(self, features):
        """
        Predicts the emotions on the given features vector.

        Parameters
        ----------
        features: list
            List of responses of the kernels at each of the face landmarks.

        Returns
        -------
        probabilities: OrderedDict
            The probabilities of each of the prototypic emotion, in format:
            {'anger': value, 'contempt': value, [...]}
        """

        # Predict the emotion probabilities on the given features
        probas = self._clf.predict_proba([features])[0]

        # Build a dictionary with the probabilities and emotion labels
        ret = OrderedDict()
        for i in range(len(self._emotions)):
            label = self._emotions[i]
            ret[label] = probas[i]

        return ret