text:
#################################################################################
# The Institute for the Design of Advanced Energy Systems Integrated Platform
# Framework (IDAES IP) was produced under the DOE Institute for the
# Design of Advanced Energy Systems (IDAES), and is copyright (c) 2018-2021
# by the software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia University
# Research Corporation, et al.  All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and
# license information.
#################################################################################
"""
Test for Caprese's module for NMPC.
"""

import pytest
from pyomo.environ import Var
from pyomo.common.collections import ComponentSet
from pyomo.dae.flatten import flatten_dae_components

from idaes.apps.caprese.categorize import (
        categorize_dae_variables,
        )
from idaes.apps.caprese.common.config import VariableCategory as VC
from idaes.apps.caprese.examples.cstr_model import make_model

__author__ = "Robert Parker"

@pytest.mark.unit
def test_categorize_1():
    """
    This test categorizes the enzyme cstr "as intended."
    Volume is used as a measurement, and solvent flow rates are
    fixed, but otherwise equations are as expected.
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_accumulation[0, 'aq', :]),
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            model.fs.cstr.control_volume.volume[0],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_meas_set = ComponentSet(init_meas_list)
    # Solvent holdup is not a measurement; we measure volume instead
    init_meas_set.remove(
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'Solvent'])

    init_alg_list = [
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0

@pytest.mark.unit
def test_categorize_2():
    """
    In this instance, temperature is "measured" (used as an initial condition)
    instead of energy_holdup, conc[P] is measured instead of holdup[P],
    and accumulation[C] is measured instead of holdup[C].
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time
    t0 = time.first()

    material_holdup = model.fs.cstr.control_volume.material_holdup
    material_accumulation = model.fs.cstr.control_volume.material_accumulation
    energy_holdup = model.fs.cstr.control_volume.energy_holdup
    energy_accumulation = model.fs.cstr.control_volume.energy_accumulation
    conc_mol = model.fs.cstr.outlet.conc_mol

    # Specify temperature instead of energy holdup
    energy_holdup[t0, 'aq'].unfix()
    model.fs.cstr.outlet.temperature[t0].fix(300)

    # Specify C_P instead of holdup
    material_holdup[t0, 'aq', 'P'].unfix()
    conc_mol[t0, 'P'].fix(0.)

    # Specify accumulation of C instead of holdup
    # You might want to do this if you want to start at steady state,
    # but don't know the value of every variable at the steady state
    # you want to start at...
    material_holdup[t0, 'aq', 'C'].unfix()
    material_accumulation[t0, 'aq', 'C'].fix(0.)

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_accumulation[0, 'aq', :]),
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.outlet.temperature[0],
            model.fs.cstr.control_volume.volume[0],
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'S'],
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'E'],
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'S'],
            model.fs.cstr.control_volume.material_accumulation[0, 'aq', 'C'],
            model.fs.cstr.outlet.conc_mol[0,'P'],
            ]
    init_meas_set = ComponentSet(init_meas_list)
    # No need to remove solvent holdup here as it was not added to this list.

    init_alg_list = [
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0

@pytest.mark.unit
def test_categorize_3():
    """
    In this test we fix one of the differential variables.
    This is the case where somebody wants to run an isothermal
    CSTR.
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    CV = model.fs.cstr.control_volume
    CV.energy_holdup.fix(300)
    CV.energy_accumulation_disc_eq.deactivate()

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            *list(model.fs.cstr.control_volume.material_accumulation[0, 'aq', :]),
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            # Energy holdup has been fixed
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.control_volume.volume[0],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_meas_set = ComponentSet(init_meas_list)
    # Solvent holdup is not a measurement; we measure volume instead
    init_meas_set.remove(
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'Solvent'])

    init_alg_list = [
        # Since energy_holdup is fixed, energy_accumulation is not
        # considered a derivative. Instead it is algebraic.
        model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0

@pytest.mark.unit
def test_categorize_4():
    """
    This tests categorization when a psuedo-steady state
    approximation is used. Energy accumulation and accumulation of E
    are fixed, the corresponding initial conditions unfixed, and
    the corresponding discretization equations deactivated.
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    CV = model.fs.cstr.control_volume
    CV.energy_accumulation[:, 'aq'].fix(0.)
    CV.material_accumulation[:, 'aq', 'E'].fix(0.)
    CV.energy_holdup[0, 'aq'].unfix()
    CV.material_holdup[0, 'aq', 'E'].unfix()
    CV.energy_accumulation_disc_eq.deactivate()
    CV.material_accumulation_disc_eq.deactivate()

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            CV.material_accumulation[0, 'aq', 'C'],
            CV.material_accumulation[0, 'aq', 'S'],
            CV.material_accumulation[0, 'aq', 'P'],
            CV.material_accumulation[0, 'aq', 'Solvent'],
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            CV.material_holdup[0, 'aq', 'C'],
            CV.material_holdup[0, 'aq', 'S'],
            CV.material_holdup[0, 'aq', 'P'],
            CV.material_holdup[0, 'aq', 'Solvent'],
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
            CV.material_accumulation[0, 'aq', 'E'],
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.control_volume.volume[0],
            CV.material_holdup[0,'aq','P'],
            CV.material_holdup[0,'aq','C'],
            CV.material_holdup[0,'aq','S'],
            ]
    init_meas_set = ComponentSet(init_meas_list)

    init_alg_list = [
        model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
        CV.material_holdup[0, 'aq', 'E'],
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0

@pytest.mark.unit
def test_categorize_error():
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)

    # Add a dummy var to treat as an input.
    # This var is not in `dae_vars`, so it will not be located during
    # categorization, which should fail.
    model.dummy_var = Var(time)

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            model.dummy_var[0],
            ]

    with pytest.raises(RuntimeError, match=r'Not all inputs could be found'):
        category_dict = categorize_dae_variables(
                dae_vars,
                time,
                init_input_list,
                )

    # Re-run flattener. Now `dummy_var` should be included in `dae_vars`.
    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(
            dae_vars,
            time,
            init_input_list,
            )

if __name__ == '__main__':
    test_categorize_1()
    test_categorize_2()
    test_categorize_3()
    test_categorize_4()
    test_categorize_error()

max_stars_repo_name: OOAmusat/idaes-pse

max_issues_repo_name: OOAmusat/idaes-pse

max_forks_repo_name: OOAmusat/idaes-pse

reason: from pyomo


text:
# Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
#
# 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 importlib.resources as pkg_resources
import logging
import os
from collections import defaultdict
from typing import List, Sequence

import gltf
import numpy
from direct.showbase.ShowBase import ShowBase
from envision import types as envision_types
from envision.client import Client as EnvisionClient
from panda3d.core import ClockObject, NodePath, Shader, loadPrcFileData

import warnings

with warnings.catch_warnings():
    # XXX: Benign warning, seems no other way to "properly" fix
    warnings.filterwarnings("ignore", "numpy.ufunc size changed")
    from sklearn.metrics.pairwise import euclidean_distances

from . import glsl, models
from .agent_manager import AgentManager
from .bubble_manager import BubbleManager
from .colors import SceneColors
from .controllers import ActionSpaceType, Controllers
from .masks import RenderMasks
from .motion_planner_provider import MotionPlannerProvider
from .provider import ProviderState
from .scenario import Scenario
from .sensors import Collision
from .sumo_road_network import SumoRoadNetwork
from .sumo_traffic_simulation import SumoTrafficSimulation
from .traffic_history_provider import TrafficHistoryProvider
from smarts.core.chassis import AckermannChassis, BoxChassis
from .trap_manager import TrapManager
from .utils import pybullet
from .utils.pybullet import bullet_client as bc
from .utils.visdom_client import VisdomClient
from .vehicle import VehicleState
from .vehicle_index import VehicleIndex
from .waypoints import Waypoints

# disable vsync otherwise we are limited to refresh-rate of screen
loadPrcFileData("", "sync-video false")
loadPrcFileData("", "model-path %s" % os.getcwd())
loadPrcFileData("", "audio-library-name null")
loadPrcFileData("", "gl-version 3 3")
loadPrcFileData("", "notify-level error")
loadPrcFileData("", "print-pipe-types false")

# https://www.panda3d.org/manual/?title=Multithreaded_Render_Pipeline
# loadPrcFileData('', 'threading-model Cull/Draw')

# have makeTextureBuffer create a visible window
# loadPrcFileData('', 'show-buffers true')

class SMARTSNotSetupError(Exception):
    pass

class SMARTS(ShowBase):
    def __init__(
        self,
        agent_interfaces,
        traffic_sim,
        envision: EnvisionClient = None,
        visdom: VisdomClient = None,
        timestep_sec=0.1,
        reset_agents_only=False,
        zoo_addrs=None,
    ):
        try:
            super().__init__(self, windowType="offscreen")
        except Exception as e:
            # Known reasons for this failing:
            raise Exception(
                "SMARTS: Error in initializing framework for opening graphical display and creating scene graph. "
                "A typical reason is display not found. Try running with different configurations of "
                "`export DISPLAY=` using `:0`, `:1`... . If this does not work please consult "
                "the documentation."
            ) from e

        gltf.patch_loader(self.loader)

        self._log = logging.getLogger(self.__class__.__name__)

        self._is_setup = False
        self._scenario: Scenario = None
        self._envision: EnvisionClient = envision
        self._visdom: VisdomClient = visdom
        self._timestep_sec = timestep_sec
        self._traffic_sim = traffic_sim
        self._motion_planner_provider = MotionPlannerProvider()
        self._traffic_history_provider = TrafficHistoryProvider()
        self._providers = [
            self._traffic_sim,
            self._motion_planner_provider,
            self._traffic_history_provider,
        ]

        # We buffer provider state between steps to compensate for TRACI's timestep delay
        self._last_provider_state = None
        self._reset_agents_only = reset_agents_only  # a.k.a "teleportation"
        self._imitation_learning_mode = False

        # Global clock always proceeds by a fixed dt on each tick
        self.taskMgr.clock.setMode(ClockObject.M_non_real_time)
        self.taskMgr.clock.setDt(timestep_sec)
        self._elapsed_sim_time = 0

        self.setBackgroundColor(0, 0, 0, 1)

        # Displayed framerate is misleading since we are not using a realtime clock
        self.setFrameRateMeter(False)

        # For macOS GUI. See our `BulletClient` docstring for details.
        # from .utils.bullet import BulletClient
        # self._bullet_client = BulletClient(pybullet.GUI)
        self._bullet_client = bc.BulletClient(pybullet.DIRECT)
        self._dynamic_action_spaces = {
            ActionSpaceType.Continuous,
            ActionSpaceType.Lane,
            ActionSpaceType.ActuatorDynamic,
            ActionSpaceType.LaneWithContinuousSpeed,
            ActionSpaceType.Trajectory,
            ActionSpaceType.MPC,
        }

        # Set up indices
        self._agent_manager = AgentManager(agent_interfaces, zoo_addrs)
        self._vehicle_index = VehicleIndex()

        # TODO: Should not be stored in SMARTS
        self._vehicle_collisions = defaultdict(list)  # list of `Collision` instances
        self._vehicle_states = []

        self._bubble_manager = None
        self._trap_manager: TrapManager = None

        # SceneGraph-related setup
        self._root_np = None
        self._vehicles_np = None
        self._road_network_np = None
        self._ground_bullet_id = None

    def step(self, agent_actions):
        if not self._is_setup:
            raise SMARTSNotSetupError("Must call reset() or setup() before stepping.")

        try:
            return self._step(agent_actions)
        except (KeyboardInterrupt, SystemExit):
            # ensure we clean-up if the user exits the simulation
            self._log.info("Simulation was interrupted by the user.")
            self.destroy()
            raise  # re-raise the KeyboardInterrupt
        except Exception as e:
            self._log.error(
                "Simulation crashed with exception. Attempting to cleanly shutdown."
            )
            self._log.exception(e)
            self.destroy()
            raise  # re-raise

    def _check_if_acting_on_active_agents(self, agent_actions):
        for agent_id in agent_actions.keys():
            if agent_id not in self._agent_manager.ego_agent_ids:
                self._log.warning(
                    f"Attempted to perform actions on non-existing agent, {agent_id} "
                )

    def _step(self, agent_actions):
        """Steps through the simulation while applying the given agent actions.
        Returns the observations, rewards, and done signals.
        """

        # Due to a limitation of our traffic simulator(SUMO) interface(TRACI), we can
        # only observe traffic state of the previous simulation step.
        #
        # To compensate for this, we:
        #
        # 1. Fetch social agent actions
        # 2. Step all providers and harmonize state
        # 3. Step bubble manager
        # 4. Calculate observation and reward
        # 5. Send observations to social agents
        # 6. Clear done agents
        # 7. Perform visualization
        #
        # In this way, observations and reward are computed with data that is
        # consistently with one step of latencey and the agent will observe consistent
        # data.
        dt = self.taskMgr.clock.get_dt()
        self._elapsed_sim_time = self.taskMgr.clock.get_frame_time()

        # 1. Fetch agent actions
        all_agent_actions = self._agent_manager.fetch_agent_actions(self, agent_actions)

        # 2. Step all providers and harmonize state
        provider_state = self._step_providers(all_agent_actions, dt)
        self._check_if_acting_on_active_agents(agent_actions)

        # 3. Step bubble manager and trap manager
        self._vehicle_index.sync()
        self._bubble_manager.step(self)
        self._trap_manager.step(self)

        # 4. Calculate observation and reward
        # We pre-compute vehicle_states here because we *think* the users will
        # want these during their observation/reward computations.
        # This is a hack to give us some short term perf wins. Longer term we
        # need to expose better support for batched computations
        self._vehicle_states = [v.state for v in self._vehicle_index.vehicles]

        # Agents
        self._agent_manager.step_sensors(self)

        # Panda3D
        # runs through the render pipeline
        # MUST perform this after step_sensors() above, and before observe() below,
        # so that all updates are ready before rendering happens per frame
        self.taskMgr.mgr.poll()

        observations, rewards, scores, dones = self._agent_manager.observe(self)

        response_for_ego = self._agent_manager.filter_response_for_ego(
            (observations, rewards, scores, dones)
        )

        # 5. Send observations to social agents
        self._agent_manager.send_observations_to_social_agents(observations)

        # 6. Clear done agents
        self._teardown_done_agents_and_vehicles(dones)

        # 7. Perform visualization
        self._try_emit_envision_state(provider_state, observations, scores)
        self._try_emit_visdom_obs(observations)

        observations, rewards, scores, dones = response_for_ego
        extras = dict(scores=scores)
        return observations, rewards, dones, extras

    def _teardown_done_agents_and_vehicles(self, dones):
        def done_vehicle_ids(dones):
            vehicle_ids = set()
            for agent_id, done in dones.items():
                if self._agent_manager.is_boid_agent(agent_id):
                    vehicle_ids.update(id_ for id_ in done if done[id_])
                elif done:
                    ids = self._vehicle_index.vehicle_ids_by_actor_id(agent_id)
                    # 0 if shadowing, 1 if active
                    assert len(ids) <= 1, f"{len(ids)} <= 1"
                    vehicle_ids.update(ids)

            return vehicle_ids

        def done_agent_ids(dones):
            agent_ids = set()
            for agent_id, done in dones.items():
                if self._agent_manager.is_boid_agent(agent_id):
                    if not self.agent_manager.is_boid_keep_alive_agent(
                        agent_id
                    ) and all(dones[agent_id].values()):
                        agent_ids.add(agent_id)
                elif done:
                    agent_ids.add(agent_id)

            return agent_ids

        # XXX: These can not be put inline because we do queries that must proceed
        #      the actual teardown.
        vehicles_to_teardown = done_vehicle_ids(dones)
        agents_to_teardown = done_agent_ids(dones)

        self._agent_manager.teardown_ego_agents(agents_to_teardown)
        self._agent_manager.teardown_social_agents(agents_to_teardown)
        self._teardown_vehicles(vehicles_to_teardown)

    def reset(self, scenario: Scenario):
        if scenario == self._scenario and self._reset_agents_only:
            vehicle_ids_to_teardown = []
            agent_ids = self._agent_manager.teardown_ego_agents()
            for agent_id in agent_ids:
                ids = self._vehicle_index.vehicle_ids_by_actor_id(agent_id)
                vehicle_ids_to_teardown.extend(ids)
            self._teardown_vehicles(set(vehicle_ids_to_teardown))
            self._trap_manager.init_traps(
                scenario.road_network, scenario.waypoints, scenario.missions
            )
            self._agent_manager.init_ego_agents(self)
            self._sync_panda3d()
        else:
            self.teardown()
            self.setup(scenario)

        # Tell history provide to ignore vehicles if we have assigned mission to them
        self._traffic_history_provider.set_replaced_ids(scenario.missions.keys())

        self.taskMgr.clock.reset()
        self._elapsed_sim_time = 0

        self._vehicle_states = [v.state for v in self._vehicle_index.vehicles]
        observations, _, _, _ = self._agent_manager.observe(self)
        observations_for_ego = self._agent_manager.reset_agents(observations)

        # Visualization
        self._try_emit_visdom_obs(observations)
        if len(self._agent_manager.ego_agent_ids):
            while len(observations_for_ego) < 1:
                observations_for_ego, _, _, _ = self.step({})

        self._reset_providers()

        return observations_for_ego

    def setup(self, scenario: Scenario):
        self._scenario = scenario

        self._root_np = NodePath("sim")
        self._root_np.reparentTo(self.render)

        with pkg_resources.path(
            glsl, "unlit_shader.vert"
        ) as vshader_path, pkg_resources.path(
            glsl, "unlit_shader.frag"
        ) as fshader_path:
            unlit_shader = Shader.load(
                Shader.SL_GLSL,
                vertex=str(vshader_path.absolute()),
                fragment=str(fshader_path.absolute()),
            )
        self._root_np.setShader(unlit_shader)
        self._setup_road_network()
        self._vehicles_np = self._root_np.attachNewNode("vehicles")

        self._bubble_manager = BubbleManager(scenario.bubbles, scenario.road_network)
        self._trap_manager = TrapManager(scenario)

        self._setup_bullet_client(self._bullet_client)
        provider_state = self._setup_providers(self._scenario)
        self._agent_manager.setup_agents(self)

        self._harmonize_providers(provider_state)
        self._last_provider_state = provider_state

        self._is_setup = True

    def add_provider(self, provider):
        self._providers.append(provider)

    def switch_ego_agent(self, agent_interface):
        self._agent_manager.switch_initial_agent(agent_interface)
        self._is_setup = False

    def _setup_road_network(self):
        glb_path = self.scenario.map_glb_filepath
        if self._road_network_np:
            self._log.debug(
                "road_network={} already exists. Removing and adding a new "
                "one from glb_path={}".format(self._road_network_np, glb_path)
            )
        model_np = self.loader.loadModel(glb_path, noCache=True)

        np = self._root_np.attachNewNode("road_network")
        model_np.reparent_to(np)
        np.hide(RenderMasks.OCCUPANCY_HIDE)
        np.setColor(SceneColors.Road.value)
        self._road_network_np = np

    def _setup_bullet_client(self, client: bc.BulletClient):
        client.resetSimulation()
        client.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)

        # PyBullet defaults the timestep to 240Hz. Several parameters are tuned with
        # this value in mind. For example the number of solver iterations and the error
        # reduction parameters (erp) for contact, friction and non-contact joints.
        # Attempting to get around this we set the number of substeps so that
        # timestep * substeps = 240Hz. Bullet (C++) does something to this effect as
        # well (https://git.io/Jvf0M), but PyBullet does not expose it.
        client.setPhysicsEngineParameter(
            fixedTimeStep=self._timestep_sec,
            numSubSteps=int(self._timestep_sec * 240),
            numSolverIterations=10,
            solverResidualThreshold=0.001,
            # warmStartingFactor=0.99
        )

        client.setGravity(0, 0, -9.8)

        plane_path = self._scenario.plane_filepath

        # 1e6 is the default value for plane length and width.
        plane_scale = (
            max(self._scenario.map_bounding_box[0], self._scenario.map_bounding_box[1])
            / 1e6
        )
        if not os.path.exists(plane_path):
            with pkg_resources.path(models, "plane.urdf") as path:
                plane_path = str(path.absolute())

        self._ground_bullet_id = client.loadURDF(
            plane_path,
            useFixedBase=True,
            basePosition=self._scenario.map_bounding_box[2],
            globalScaling=1.1 * plane_scale,
        )

    def teardown(self):
        self._agent_manager.teardown()
        self._vehicle_index.teardown()

        self._bullet_client.resetSimulation()
        self._traffic_sim.teardown()
        self._teardown_providers()

        if self._root_np is not None:
            self._root_np.clearLight()
            self._root_np.removeNode()

        if self._bubble_manager is not None:
            self._bubble_manager.teardown()
            self._bubble_manager = None

        if self._trap_manager is not None:
            self._trap_manager.teardown()
            self._trap_manager = None

        self._vehicles_np = None
        self._ground_bullet_id = None
        self._road_network_np = None
        self._is_setup = False

    def destroy(self):
        self.teardown()

        if self._envision:
            self._envision.teardown()

        if self._visdom:
            self._visdom.teardown()

        self._agent_manager.destroy()
        self._traffic_sim.destroy()
        self._bullet_client.disconnect()

        super().destroy()

    def _teardown_vehicles(self, vehicle_ids):
        self._vehicle_index.teardown_vehicles_by_vehicle_ids(vehicle_ids)
        self._clear_collisions(vehicle_ids)

    def attach_sensors_to_vehicles(self, agent_spec, vehicle_ids):
        self._agent_manager.attach_sensors_to_vehicles(
            self, agent_spec.interface, vehicle_ids
        )

    def observe_from(self, vehicle_ids):
        return self._agent_manager.observe_from(self, vehicle_ids)

    @property
    def road_stiffness(self):
        return self._bullet_client.getDynamicsInfo(self._ground_bullet_id, -1)[9]

    @property
    def dynamic_action_spaces(self):
        return self._dynamic_action_spaces

    @property
    def traffic_sim(self) -> SumoTrafficSimulation:
        return self._traffic_sim

    @property
    def waypoints(self) -> Waypoints:
        return self.scenario.waypoints

    @property
    def road_network(self) -> SumoRoadNetwork:
        return self.scenario.road_network

    @property
    def np(self):
        return self._root_np

    @property
    def vehicles_np(self):
        return self._vehicles_np

    @property
    def bc(self):
        return self._bullet_client

    @property
    def envision(self):
        return self._envision

    @property
    def elapsed_sim_time(self):
        return self._elapsed_sim_time

    def teardown_agents_without_vehicles(self, agent_ids: Sequence):
        """
        Teardown agents in the given list that have no vehicles registered as
        controlled-by or shadowed-by

        Params:
            agent_ids: Sequence of agent ids
        """
        agents_to_teardown = {
            agent_id
            for agent_id in agent_ids
            # Only clean-up when there are no controlled agents left (e.g. boids)
            if len(
                self._vehicle_index.vehicles_by_actor_id(
                    agent_id, include_shadowers=True
                )
            )
            == 0
        }

        agents_to_teardown = {
            id_
            for id_ in agents_to_teardown
            if not self.agent_manager.is_boid_keep_alive_agent(id_)
        }
        self.agent_manager.teardown_social_agents(filter_ids=agents_to_teardown)

    def _teardown_vehicles_and_agents(self, vehicle_ids):
        shadow_and_controlling_agents = set()
        for vehicle_id in vehicle_ids:
            agent_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle_id)
            if agent_id:
                shadow_and_controlling_agents.add(agent_id)

            shadow_agent_id = self._vehicle_index.shadow_actor_id_from_vehicle_id(
                vehicle_id
            )
            if shadow_agent_id:
                shadow_and_controlling_agents.add(shadow_agent_id)

        self._vehicle_index.teardown_vehicles_by_vehicle_ids(vehicle_ids)
        self.teardown_agents_without_vehicles(shadow_and_controlling_agents)

    def _pybullet_provider_sync(self, provider_state: ProviderState):
        current_vehicle_ids = {v.vehicle_id for v in provider_state.vehicles}
        previous_sv_ids = self._vehicle_index.social_vehicle_ids()
        exited_vehicles = previous_sv_ids - current_vehicle_ids
        self._teardown_vehicles_and_agents(exited_vehicles)

        # Update our pybullet world given this provider state
        for vehicle in provider_state.vehicles:
            vehicle_id = vehicle.vehicle_id
            # either this is a pybullet agent vehicle, or it is a social vehicle
            if vehicle_id in self._vehicle_index.agent_vehicle_ids():
                # this is an agent vehicle
                agent_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle_id)
                agent_interface = self._agent_manager.agent_interface_for_agent_id(
                    agent_id
                )
                agent_action_space = agent_interface.action_space
                if agent_action_space not in self._dynamic_action_spaces:
                    # This is not a pybullet agent, but it has an avatar in this world
                    # to make it's observations. Update the avatar to match the new
                    # state of this vehicle
                    pybullet_vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
                    assert isinstance(pybullet_vehicle.chassis, BoxChassis)
                    pybullet_vehicle.control(pose=vehicle.pose, speed=vehicle.speed)
            else:
                # This vehicle is a social vehicle
                if vehicle_id in self._vehicle_index.social_vehicle_ids():
                    social_vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
                else:
                    # It is a new social vehicle we have not seen yet.
                    # Create it's avatar.
                    social_vehicle = self._vehicle_index.build_social_vehicle(
                        sim=self,
                        vehicle_state=vehicle,
                        actor_id=vehicle_id,
                        vehicle_id=vehicle_id,
                        vehicle_type=vehicle.vehicle_type,
                    )
                # Update the social vehicle avatar to match the vehicle state
                social_vehicle.control(pose=vehicle.pose, speed=vehicle.speed)

    def _pybullet_provider_step(self, agent_actions) -> ProviderState:
        self._perform_agent_actions(agent_actions)

        self._bullet_client.stepSimulation()

        self._process_collisions()

        provider_state = ProviderState()
        pybullet_agent_ids = {
            agent_id
            for agent_id, interface in self._agent_manager.agent_interfaces.items()
            if interface.action_space in self._dynamic_action_spaces
        }

        for vehicle_id in self._vehicle_index.agent_vehicle_ids():
            agent_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle_id)
            if agent_id not in pybullet_agent_ids:
                continue

            vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
            vehicle.step(self._elapsed_sim_time)
            provider_state.vehicles.append(
                VehicleState(
                    vehicle_id=vehicle.id,
                    vehicle_type="passenger",
                    pose=vehicle.pose,
                    dimensions=vehicle.chassis.dimensions,
                    speed=vehicle.speed,
                    source="PYBULLET",
                )
            )

        return provider_state

    @property
    def vehicle_index(self):
        return self._vehicle_index

    @property
    def agent_manager(self):
        return self._agent_manager

    @property
    def providers(self):
        # TODO: Add check to ensure that action spaces are disjoint between providers
        # TODO: It's inconsistent that pybullet is not here
        return self._providers

    def _setup_providers(self, scenario) -> ProviderState:
        provider_state = ProviderState()
        for provider in self.providers:
            provider_state.merge(provider.setup(scenario))

        return provider_state

    def _teardown_providers(self):
        for provider in self.providers:
            provider.teardown()

        self._last_provider_state = None

    def _harmonize_providers(self, provider_state: ProviderState):
        for provider in self.providers:
            provider.sync(provider_state)

        self._pybullet_provider_sync(provider_state)
        self._sync_panda3d()

    def _reset_providers(self):
        for provider in self.providers:
            provider.reset()

    def _step_providers(self, actions, dt) -> List[VehicleState]:
        accumulated_provider_state = ProviderState()

        def agent_controls_vehicles(agent_id):
            vehicles = self._vehicle_index.vehicles_by_actor_id(agent_id)
            return len(vehicles) > 0

        def matches_provider_action_spaces(agent_id, action_spaces):
            interface = self._agent_manager.agent_interface_for_agent_id(agent_id)
            return interface.action_space in action_spaces

        # PyBullet
        pybullet_actions = {
            agent_id: action
            for agent_id, action in actions.items()
            if agent_controls_vehicles(agent_id)
            and matches_provider_action_spaces(agent_id, self._dynamic_action_spaces)
        }
        accumulated_provider_state.merge(self._pybullet_provider_step(pybullet_actions))

        for provider in self.providers:
            provider_state = self._step_provider(
                provider, actions, dt, self._elapsed_sim_time
            )
            if provider == self._traffic_sim:
                # Remove agent vehicles from provider vehicles
                provider_state.filter(self._vehicle_index.agent_vehicle_ids())

            accumulated_provider_state.merge(provider_state)

        self._harmonize_providers(accumulated_provider_state)
        return accumulated_provider_state

    def _step_provider(self, provider, actions, dt, elapsed_sim_time):
        def agent_controls_vehicles(agent_id):
            vehicles = self._vehicle_index.vehicles_by_actor_id(agent_id)
            return len(vehicles) > 0

        provider_actions = {}
        for agent_id, action in actions.items():
            agent_interface = self._agent_manager.agent_interface_for_agent_id(agent_id)
            if (
                agent_interface
                and agent_controls_vehicles(agent_id)
                and agent_interface.action_space in provider.action_spaces
            ):
                vehicle_ids = [
                    v.id
                    for v in self._vehicle_index.vehicles_by_actor_id(
                        agent_id, include_shadowers=True
                    )
                ]

                if self._agent_manager.is_boid_agent(agent_id):
                    for vehicle_id, vehicle_action in action.items():
                        assert vehicle_id in vehicle_ids
                        provider_actions[vehicle_id] = vehicle_action
                else:
                    assert len(vehicle_ids) == 1
                    provider_actions[vehicle_ids[0]] = action

        provider_state = provider.step(provider_actions, dt, elapsed_sim_time)
        return provider_state

    @property
    def scenario(self):
        return self._scenario

    @property
    def traffic_sim(self):
        return self._traffic_sim

    @property
    def np(self):
        return self._root_np

    @property
    def timestep_sec(self):
        return self._timestep_sec

    @property
    def road_stiffness(self):
        return self._bullet_client.getDynamicsInfo(self._ground_bullet_id, -1)[9]

    def neighborhood_vehicles_around_vehicle(self, vehicle, radius=None):
        other_states = [v for v in self._vehicle_states if v.vehicle_id != vehicle.id]
        if radius is None:
            return other_states

        other_positions = [state.pose.position for state in other_states]
        if not other_positions:
            return []

        distances = euclidean_distances(other_positions, [vehicle.position]).reshape(
            -1,
        )
        indices = numpy.argwhere(distances <= radius).flatten()
        return [other_states[i] for i in indices]

    def vehicle_did_collide(self, vehicle_id):
        return (
            len(
                [
                    c
                    for c in self._vehicle_collisions[vehicle_id]
                    if c.collidee_id != self._ground_bullet_id
                ]
            )
            > 0
        )

    def vehicle_collisions(self, vehicle_id):
        return [
            c
            for c in self._vehicle_collisions[vehicle_id]
            if c.collidee_id != self._ground_bullet_id
        ]

    def _clear_collisions(self, vehicle_ids):
        for vehicle_id in vehicle_ids:
            self._vehicle_collisions.pop(vehicle_id, None)

    def _perform_agent_actions(self, agent_actions):
        for agent_id, action in agent_actions.items():
            agent_vehicles = self._vehicle_index.vehicles_by_actor_id(agent_id)
            if len(agent_vehicles) == 0:
                self._log.warning(
                    f"{agent_id} doesn't have a vehicle, is the agent done? (dropping action)"
                )
            else:
                agent_interface = self._agent_manager.agent_interface_for_agent_id(
                    agent_id
                )
                is_boid_agent = self._agent_manager.is_boid_agent(agent_id)

                for vehicle in agent_vehicles:
                    vehicle_action = action[vehicle.id] if is_boid_agent else action

                    controller_state = self._vehicle_index.controller_state_for_vehicle_id(
                        vehicle.id
                    )
                    sensor_state = self._vehicle_index.sensor_state_for_vehicle_id(
                        vehicle.id
                    )
                    # TODO: Support performing batched actions
                    Controllers.perform_action(
                        self,
                        agent_id,
                        vehicle,
                        vehicle_action,
                        controller_state,
                        sensor_state,
                        agent_interface.action_space,
                        agent_interface.vehicle_type,
                    )

    def _sync_panda3d(self):
        for vehicle in self._vehicle_index.vehicles:
            vehicle.sync_to_panda3d()

    def _process_collisions(self):
        self._vehicle_collisions = defaultdict(list)  # list of `Collision` instances

        for vehicle_id in self._vehicle_index.agent_vehicle_ids():
            vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
            # We are only concerned with vehicle-vehicle collisions
            collidee_bullet_ids = set(
                [p.bullet_id for p in vehicle.chassis.contact_points]
            )
            collidee_bullet_ids.discard(self._ground_bullet_id)

            if not collidee_bullet_ids:
                continue

            for bullet_id in collidee_bullet_ids:
                collidee = self._bullet_id_to_vehicle(bullet_id)
                collision = self._node_to_collision(collidee.np.node())
                self._vehicle_collisions[vehicle_id].append(collision)

    def _bullet_id_to_vehicle(self, bullet_id):
        for vehicle in self._vehicle_index.vehicles:
            if bullet_id == vehicle.chassis.bullet_id:
                return vehicle

        assert False, "Only collisions with agent or social vehicles is supported"

    def _node_to_collision(self, node):
        for vehicle in self._vehicle_index.vehicles:
            if node == vehicle.np.node():
                actor_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle.id)
                # TODO: Should we specify the collidee as the vehicle ID instead of
                #       the agent/social ID?
                return Collision(collidee_id=actor_id)

        assert False, "Only collisions with agent or social vehicles is supported"

    def _try_emit_envision_state(self, provider_state, obs, scores):
        if not self._envision:
            return

        traffic = {}
        for v in provider_state.vehicles:
            if v.vehicle_id in self._vehicle_index.agent_vehicle_ids():
                # this is an agent controlled vehicle
                agent_id = self._vehicle_index.actor_id_from_vehicle_id(v.vehicle_id)
                agent_obs = obs[agent_id]
                is_boid_agent = self._agent_manager.is_boid_agent(agent_id)
                vehicle_obs = agent_obs[v.vehicle_id] if is_boid_agent else agent_obs

                if self._agent_manager.is_ego(agent_id):
                    actor_type = envision_types.TrafficActorType.Agent
                    mission_route_geometry = self._vehicle_index.sensor_state_for_vehicle_id(
                        v.vehicle_id
                    ).mission_planner.route.geometry
                else:
                    actor_type = envision_types.TrafficActorType.SocialAgent
                    mission_route_geometry = None

                point_cloud = vehicle_obs.lidar_point_cloud or ([], [], [])
                point_cloud = point_cloud[0]  # (points, hits, rays), just want points

                # TODO: driven path should be read from vehicle_obs
                driven_path = self._vehicle_index.vehicle_by_id(
                    v.vehicle_id
                ).driven_path_sensor()

                road_waypoints = []
                if vehicle_obs.road_waypoints:
                    road_waypoints = [
                        path
                        for paths in vehicle_obs.road_waypoints.lanes.values()
                        for path in paths
                    ]
                traffic[v.vehicle_id] = envision_types.TrafficActorState(
                    name=self._agent_manager.agent_name(agent_id),
                    actor_type=actor_type,
                    vehicle_type=envision_types.VehicleType.Car,
                    position=v.pose.position,
                    heading=v.pose.heading,
                    speed=v.speed,
                    actor_id=envision_types.format_actor_id(
                        agent_id, v.vehicle_id, is_multi=is_boid_agent,
                    ),
                    events=vehicle_obs.events,
                    waypoint_paths=(vehicle_obs.waypoint_paths or []) + road_waypoints,
                    point_cloud=point_cloud,
                    driven_path=driven_path,
                    mission_route_geometry=mission_route_geometry,
                )
            elif v.vehicle_id in self._vehicle_index.social_vehicle_ids():
                # this is a social vehicle
                traffic[v.vehicle_id] = envision_types.TrafficActorState(
                    actor_type=envision_types.TrafficActorType.SocialVehicle,
                    vehicle_type=v.vehicle_type,
                    position=list(v.pose.position),
                    heading=v.pose.heading,
                    speed=v.speed,
                )

        bubble_geometry = [
            list(bubble.geometry.exterior.coords)
            for bubble in self._bubble_manager.bubbles
        ]

        state = envision_types.State(
            traffic=traffic,
            scenario_id=self.scenario.scenario_hash,
            bubbles=bubble_geometry,
            scene_colors=SceneColors.EnvisionColors.value,
            scores=scores,
        )
        self._envision.send(state)

    def _try_emit_visdom_obs(self, obs):
        if not self._visdom:
            return

        self._visdom.send(obs)

ext: py

reason: import numpy


text:
"""
Load and parse through an input data file. Currently using "centrals-sample.fits".
Creates and populates a database table to represent centrals model
"""
import os,sys
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "legacyhalos_web.settings")
import django
django.setup()

from legacyhalos_web.models import Centrals
#import csv
from astropy.table import Table

#if True:
## csv_filepathname="centrals-sample.csv"
## dataReader = csv.reader(open(csv_filepathname), delimiter=',', quotechar='"')
## for row in dataReader:
##     print('Row:', row)
##     centrals=Centrals()
##     centrals.objid=row[0]
##     centrals.morphtype=row[1]
##     centrals.ra = row[2]
##     centrals.dec = row[3]
##     centrals.mem_match_id = row[4]
##     centrals.z = row[5]
##     centrals.la = row[6]
##     centrals.sdss_objid = row[7]
##     centrals.save()

#
data = Table.read('centrals-sample.fits')
data.rename_column('type', 'morphtype')

for row in data:
    #print('Row:', row)
    centrals=Centrals()
    centrals.objid=row[0]
    centrals.morphtype=row[1]
    centrals.ra = row[2]
    centrals.dec = row[3]
    centrals.mem_match_id_string = '%0*d' % (7, row[4])
    centrals.mem_match_id = row[4]
    centrals.z = row[5]
    centrals.la = row[6]
    centrals.sdss_objid = row[7]
    #centrals.viewer_link = centrals.viewer_link()
    #centrals.skyserver_link = centrals.skyserver_link()
    centrals.viewer_link = 'http://legacysurvey.org/viewer/?ra={:.6f}&dec={:.6f}&zoom=15&layer=decals-dr5'.format(centrals.ra, centrals.dec)
    centrals.skyserver_link = 'http://skyserver.sdss.org/dr14/en/tools/explore/summary.aspx?id={}'.format(centrals.sdss_objid)

    centrals.save()

## for row in data:
##     centrals = Centrals()
##     for col in data.colnames:
##         col.getattr(centrals, col) = data[row][col]
##     centrals.save()

max_stars_repo_name: moustakas/legacyhalos-web

max_issues_repo_name: moustakas/legacyhalos-web

max_forks_repo_name: moustakas/legacyhalos-web

reason: from astropy


text:
import numpy as np
import scipy.stats as stats

def mean_median_mode(array):
    """
    Returns the mean, median and mode of an array

    Parameters
    ----------
    array: numpy.ndarray

    Returns:
    mean, median, mode

    """
    mean = np.mean(array)
    median = np.median(array)
    mode = stats.mode(array)[0][0]

    return mean, median, mode

def gaussian(x, mu, sig, scale=None, normed=False):
    """
    Generates a gaussian curve

    Parameters:
    x : numpy.ndarray
        x values

    mu : float or int
        The mean of the gaussian curve.

    sig : float or int
        The standard deviation of the gaussian curve.

    scale : float or int

    normed : boolean, optional
        If normed is true, the gassuan returned will be normalised such
        that the area under the curve intergrates to unity.
        Setting `normed = True` overrides the value in scale.
        Default: False

    """
    if scale is None:
        a = 1

    else:
        a = scale

    if normed:
        a = 1 / (sig * np.sqrt(2 * np.pi))

    return a * np.exp(-(x - mu)**2.0 / (2 * sig**2.0))

max_stars_repo_name: abatten/pyx

max_issues_repo_name: abatten/pyx

max_forks_repo_name: abatten/pyx

reason: import numpy,import scipy


text:
import os
import pickle
import signal
from io import BytesIO
from typing import Tuple, List, Union, Dict, Sized, Iterable, Generator, Optional, Hashable

import msgpack
import numpy as np
import torch
from torch import Tensor

from lz4 import frame
from multiprocessing.reduction import ForkingPickler

# Setup fork-server mode for the entire library
# Preload common libraries to speed up loading time
multiprocessing = mp_ctx = torch.multiprocessing.get_context('forkserver')
mp_ctx.set_forkserver_preload(["numpy", "torch", "zmq", "pickle", "msgpack", "lz4"])

# Hotswap the process to a thread for debugging
# from threading import Thread
# mp_ctx.Process = Thread

# Recursive definitions not supported yet, so I use Tensor for the subtypes
# But the correct type should have BufferType where Tensor is.
BufferType = Union[Tensor, List[Tensor], Dict[Hashable, Tensor]]
ShapeBufferType = Union[Tuple[int, ...], List[Tuple[int, ...]], Dict[Hashable, Tuple[int, ...]]]
DtypeBufferType = Union[torch.dtype, List[torch.dtype], Dict[Hashable, torch.dtype]]

DEBUG = False
VERBOSE = False

def serialize_int(x: int) -> bytes:
    """ Efficiently convert a python integer to a serializable byte-string. """
    byte_size = (int.bit_length(x) + 8) // 8
    return int.to_bytes(x, length=byte_size, byteorder='big')

def deserialize_int(b: bytes) -> int:
    """ Recover integer value from bytestring. """
    return int.from_bytes(b, byteorder='big')

def serialize_tensor(tensor: Tensor) -> bytes:
    """ Serialize a tensor by placing it in shared memory. """
    # return pickle.dumps(tensor)
    buf = BytesIO()
    ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(tensor)
    return buf.getvalue()

def deserialize_tensor(serialized_tensor: bytes) -> Tensor:
    """ Convert back from a serialize tensor. Not that this can only be done once for a given bytestring. """
    return pickle.loads(serialized_tensor)

def send_sigkill(pid: int):
    """ Forcefully kill a process by PID.

    Parameters
    ----------
    pid: int
        Process to kill
    """
    os.kill(pid, signal.SIGKILL)

def relative_channel(base_channel: str, prefix: str) -> str:
    """ Convert an IPC channel into one that is prefixed by a temporary directory.

    This allows us to make multiple NetworkManagers that dont interfere with each other.

    Parameters
    ----------
    base_channel: str
        The IPC channel to start with
    prefix : str
        The prefix for the IPC path.

    Returns
    -------
    str
        The updated relative path.

    """
    protocol, path = base_channel.split("//")
    return "{}//{}/{}".format(protocol, prefix, path)

def iterate_window(iterator: Union[Sized, Iterable], n: int = 2) -> Generator:
    """ Iterate over a sliding window of an iterator with no overlap.

    Parameters
    ----------
    iterator: Iterable
        The target to iterate over. Must support __len__. Total length must be a multiple of n.
    n: int
        Window size

    Returns
    -------
    Iterable
        Generator for viewing every n elements of the iterator.
    """
    size = len(iterator)
    iterator = iter(iterator)
    for _ in range(0, size, n):
        yield (next(iterator) for _ in range(n))

def optional(variable: Optional[object], default: object):
    """ Optional parameter that is default if None.
    """
    return default if variable is None else variable

def _serialize_compressed_buffer(buffer: BufferType, compress: int):
    if isinstance(buffer, dict):
        return {key: _serialize_compressed_buffer(val, compress) for key, val in buffer.items()}
    elif isinstance(buffer, (list, tuple)):
        return [_serialize_compressed_buffer(val, compress) for val in buffer]
    elif isinstance(buffer, torch.Tensor):
        buffer = buffer.numpy()

    shape = buffer.shape
    dtype = buffer.dtype.str

    if len(shape) > 0:
        pickled = buffer.view(np.uint8).data
    else:
        pickled = buffer.tobytes()

    pickled = frame.compress(pickled, compression_level=compress)

    return dtype, shape, pickled

def _serialize_uncompressed_buffer(buffer: BufferType):
    if isinstance(buffer, dict):
        return {key: _serialize_uncompressed_buffer(val) for key, val in buffer.items()}
    elif isinstance(buffer, (list, tuple)):
        return [_serialize_uncompressed_buffer(val) for val in buffer]
    elif isinstance(buffer, torch.Tensor):
        buffer = buffer.numpy()

    shape = buffer.shape
    dtype = buffer.dtype.str

    if len(shape) > 0:
        pickled = buffer.view(np.uint8).data
    else:
        pickled = buffer.tobytes()

    return dtype, shape, pickled

def serialize_buffer(buffer: BufferType, compress: int = -1):
    if compress < 0:
        serialized = _serialize_uncompressed_buffer(buffer)
    else:
        serialized = _serialize_compressed_buffer(buffer, compress)

    serialized = msgpack.dumps((compress, serialized), use_bin_type=True)
    return serialized

def _deserialize_compressed_buffer(serialized: bytes):
    if isinstance(serialized, dict):
        return {key: _deserialize_compressed_buffer(val) for key, val in serialized.items()}
    elif not isinstance(serialized[0], str):
        return [_deserialize_compressed_buffer(val) for val in serialized]

    dtype, shape, pickled = serialized
    array = np.frombuffer(frame.decompress(pickled), np.dtype(dtype)).reshape(shape)
    return torch.from_numpy(array)

def _deserialize_uncompressed_buffer(serialized: bytes):
    if isinstance(serialized, dict):
        return {key: _deserialize_uncompressed_buffer(val) for key, val in serialized.items()}
    elif not isinstance(serialized[0], str):
        return [_deserialize_uncompressed_buffer(val) for val in serialized]

    dtype, shape, pickled = serialized
    array = np.frombuffer(pickled, np.dtype(dtype)).reshape(shape)
    return torch.from_numpy(array)

def deserialize_buffer(serialized: bytes):
    compress, serialized = msgpack.loads(serialized, raw=False, use_list=False)
    if compress < 0:
        return _deserialize_uncompressed_buffer(serialized)
    else:
        return _deserialize_compressed_buffer(serialized)

def _deserialize_compressed_buffer_into(to_buffer: BufferType, serialized: bytes):
    if isinstance(to_buffer, dict):
        for key, to_tensor in to_buffer.items():
            size = _deserialize_compressed_buffer_into(to_tensor, serialized[key])
        return size

    elif isinstance(to_buffer, (list, tuple)):
        for to_tensor, from_tensor in zip(to_buffer, serialized):
            size = _deserialize_compressed_buffer_into(to_tensor, from_tensor)
        return size

    else:
        dtype, shape, pickled = serialized
        from_buffer = np.frombuffer(frame.decompress(pickled), np.dtype(dtype)).reshape(shape)
        from_buffer = torch.from_numpy(from_buffer)

        size = shape[0]
        to_buffer[:size].copy_(from_buffer)

        return size

def _deserialize_uncompressed_buffer_into(to_buffer: BufferType, serialized: bytes):
    if isinstance(to_buffer, dict):
        for key, to_tensor in to_buffer.items():
            size = _deserialize_uncompressed_buffer_into(to_tensor, serialized[key])
        return size

    elif isinstance(to_buffer, (list, tuple)):
        for to_tensor, from_tensor in zip(to_buffer, serialized):
            size = _deserialize_uncompressed_buffer_into(to_tensor, from_tensor)
        return size

    else:
        dtype, shape, pickled = serialized
        from_buffer = np.frombuffer(pickled, np.dtype(dtype)).reshape(shape)
        from_buffer = torch.from_numpy(from_buffer)

        size = shape[0]
        to_buffer[:size].copy_(from_buffer)

        return size

def deserialize_buffer_into(to_buffer: BufferType, serialized: bytes):
    compress, serialized = msgpack.loads(serialized, raw=False, use_list=False)
    if compress < 0:
        size = _deserialize_uncompressed_buffer_into(to_buffer, serialized)
    else:
        size = _deserialize_compressed_buffer_into(to_buffer, serialized)
    return size, compress

class IdentityContextManager:
    def __call__(self, *args, **kwargs):
        return self

    def __enter__(self):
        return None

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


text:
# coding=utf-8
# Copyright 2019-present, the HuggingFace Inc. team and Facebook, 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.
""" Dataset to distilled models
    adapted in part from Facebook, Inc XLM model (https://github.com/facebookresearch/XLM)
"""
import numpy as np
import torch
import os
from torch.utils.data import Dataset

from utils import logger
import json

class LazyLmSeqsDataset(Dataset):
    """Custom Dataset wrapping language modeling sequences.

    Each sample will be retrieved by indexing the list of token_ids and their corresponding lengths.

    Input:
    ------
        params: `NameSpace` parameters
        data: `List[np.array[int]]
    """

    def __init__(self, params):
        self.params = params
        self.train_data_info = {}
        self.sampleid2fileid = []
        self.create_dataset_metadata()
        self.total_samples = len(self.sampleid2fileid)
        self.print_statistics()

    def __getitem__(self, index):
        id_info = self.sampleid2fileid[index]
        file_id = id_info["file_id"]
        local_sample_id = id_info["local_sample_id"]

        file_name = self.train_data_info["info"][file_id]["file_name"]
        data = self.memmap_read_npy(os.path.join(self.params.data_file, file_name))
        sample = data[local_sample_id]
        return (sample, sample.shape[0])

    def __len__(self):
        return self.total_samples

    def print_statistics(self):
        """
        Print some statistics on the corpus. Only the master process.
        """
        if not self.params.is_master:
            return
        logger.info(f"{len(self)} sequences")
        # data_len = sum(self.lengths)
        # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
        # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

        # unk_idx = self.params.special_tok_ids['unk_token']
        # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
        # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')

    def memmap_read_npy(self, npy_file):
        return np.lib.format.open_memmap(npy_file,
                          mode='r+', dtype=np.uint16,
                          shape=None, fortran_order=False, version=None)

    def create_dataset_metadata(self):
        train_npy_path = self.params.data_file
        if not os.path.isdir(train_npy_path): raise IOError("Please provide the folder path not file.")
        train_npy_files = os.listdir(train_npy_path)
        self.train_data_info["info"] = []
        file_id = len(self.train_data_info["info"])
        for npy_file in train_npy_files:
            if npy_file.endswith("npy"):
                data = self.memmap_read_npy(os.path.join(train_npy_path, npy_file))
                data_shape = data.shape

                info = {
                    "file_id" : file_id,
                    "file_name" : npy_file,
                    "n_samples" : data_shape[0],
                    "sample_size" : data_shape[1],
                    "dtype" : data.dtype
                }
                self.sampleid2fileid += [{"file_id":file_id, "local_sample_id":local_sample_id}
                                             for local_sample_id in range(data_shape[0])]
                self.train_data_info["info"].append(info)

#         with open("train_data_info.json", "w") as f:
#             f.write(json.dumps(self.train_data_info))

    def batch_sequences(self, batch):
        """
        Do the padding and transform into torch.tensor.
        """
        token_ids = [t[0] for t in batch]
        lengths = [t[1] for t in batch]
        assert len(token_ids) == len(lengths)

        # Max for paddings
        max_seq_len_ = max(lengths)

        # Pad token ids
        if self.params.mlm:
            pad_idx = self.params.special_tok_ids["pad_token"]
        else:
            pad_idx = self.params.special_tok_ids["unk_token"]
        tk_ = [list(t.astype(int)) + [pad_idx] * (max_seq_len_ - len(t)) for t in token_ids]
        assert len(tk_) == len(token_ids)
        assert all(len(t) == max_seq_len_ for t in tk_)

        tk_t = torch.tensor(tk_)  # (bs, max_seq_len_)
        lg_t = torch.tensor(lengths)  # (bs)
        return tk_t, lg_t

max_stars_repo_name: kinoc/distilling

max_issues_repo_name: kinoc/distilling

max_forks_repo_name: kinoc/distilling