# -*- coding: utf-8 -*-
"""
Tools for designing and determining operational parameters of a closed-end
detonation tube with optical access.
"""

import os
import warnings
from math import sqrt

import cantera as ct
import numpy as np
import pandas as pd
import pint

from . import thermochem, units
from .thermochem import _U

_Q = _U.Quantity

_DIR_LOOKUP_DATA = os.path.join(os.path.dirname(os.path.relpath(__file__)), "lookup_data")
MATERIAL_LIMITS = dict(
    welded=pd.read_csv(
        os.path.join(_DIR_LOOKUP_DATA, "ASME_B31_1_stress_limits_welded.csv"),
        # index_col=0
    ),
    seamless=pd.read_csv(
        os.path.join(_DIR_LOOKUP_DATA, "ASME_B31_1_stress_limits_seamless.csv"),
        # index_col=0
    ),
)
PIPE_SCHEDULES = pd.read_csv(os.path.relpath(os.path.join(_DIR_LOOKUP_DATA, "pipe_schedules.csv")), index_col=0)

def _collect_tube_materials():
    """
    Reads in a csv file containing tube materials, their corresponding
    ASME B16.5 material groups, and selected material properties.

    Returns
    -------
    df_materials : pd.DataFrame
        Dataframe of materials and their corresponding material groups and
        properties
    """
    file_directory = os.path.join(os.path.dirname(os.path.relpath(__file__)), "lookup_data")
    file_name = "materials_list.csv"
    file_location = os.path.relpath(os.path.join(file_directory, file_name))

    # read in csv and extract information
    if os.path.exists(file_location):
        # noinspection PyUnresolvedReferences
        try:
            df_materials = pd.read_csv(file_location)
            # type: pd.DataFrame
        except pd.errors.EmptyDataError:
            raise ValueError("\n" + file_name + " is empty")

    else:
        # raise an exception if the file doesn't exist
        raise ValueError("\n" + file_name + " does not exist")

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        # apply units
        df_materials["ElasticModulus"] = [_Q(item, "GPa") for item in df_materials["ElasticModulus"].values]
        df_materials["Density"] = [_Q(item, "g/cm^3") for item in df_materials["Density"].values]

    return df_materials

TUBE_MATERIALS = _collect_tube_materials()

def _check_materials_list():
    """
    Makes sure that the materials in materials_list.csv have stress limits
    and flange ratings. This function relies on _get_material_groups(), and
    either raises an error or returns True.

    Returns
    -------
    True
    """
    # collect files
    file_directory = os.path.join(os.path.dirname(os.path.relpath(__file__)), "lookup_data")
    my_files = os.listdir(file_directory)
    flange_ratings = [file for file in my_files if "flange" in file.lower()]
    stress_limits = [file for file in my_files if "stress" in file.lower()]

    # make sure things were actually loaded
    if not bool(flange_ratings + stress_limits):
        raise FileNotFoundError("No files containing 'flange' or 'stress' found")

    # initialize an error string and error indicator. Error string will be
    # used to aggregate errors in the list of available materials so that
    # all issues may be rectified simultaneously.
    error_string = "\n"
    has_errors = False

    # make sure all pipe material limits are either welded or seamless
    # other types are permitted, but will raise a warning
    for file in stress_limits:
        if ("welded" not in file.lower()) and ("seamless" not in file.lower()):
            # warn that something is weird
            warnings.warn(file + " does not indicate whether it is welded or seamless")

        # check the first row of the file in question to extract the names
        # of the materials that it contains stress limits for
        file_location = os.path.join(file_directory, file)
        with open(file_location, "r") as current_file:
            # read the first line, strip off carriage return, and split by
            # comma separators. Ignore first value, as this is temperature.
            materials = current_file.readline().strip().split(",")[1:]

            # check to make sure that each material in the list of available
            # materials has a stress limit curve for the current limit type
            for item in TUBE_MATERIALS["Grade"].values.astype(str):
                if item not in materials:
                    # a material is missing from the limits spreadsheet.
                    # indicate that an error has occurred, and add it to the
                    # error string.
                    error_string += "Material " + item + " not found in " + file_location + "\n"
                    has_errors = True

    # find out which material groups need to be inspected
    groups = set()
    for group in TUBE_MATERIALS["Group"].values.astype(str):
        groups.add(group.replace(".", "_"))

    # check folder to make sure the correct files exist
    for group in groups:
        if not any(rating.find(group) > 0 for rating in flange_ratings):
            # current group was not found in any of the files
            error_string += "material group " + group + " not found" + "\n"
            has_errors = True

    # report all errors
    if has_errors:
        raise ValueError(error_string)

    return True

_check_materials_list()

def _check_material(material):
    if material not in TUBE_MATERIALS.Grade.values:
        msg = "%s not found in materials." " Try Flange.available_materials" % material
        raise ValueError(msg)
    return True

def _get_flange_limits_from_csv():
    """
    Reads in flange pressure limits as a function of temperature for
    different pressure classes per ASME B16.5. Temperature is in Centigrade
    and pressure is in bar.
    """
    groups = ["2.1", "2.2", "2.3"]
    group_flange_limits = {group: None for group in groups}

    for group in groups:
        # ensure group is valid
        file_group = str(group).replace(".", "_")
        file_directory = os.path.join(os.path.dirname(os.path.relpath(__file__)), "lookup_data")
        file_name = "ASME_B16_5_flange_ratings_group_" + file_group + ".csv"
        file_location = os.path.relpath(os.path.join(file_directory, file_name))
        if not os.path.exists(file_location):
            raise FileNotFoundError("\n" + file_location + "not found")

        # import the correct .csv file as a pandas dataframe
        flange_limits = pd.read_csv(file_location)

        # ensure all temperatures and pressures are floats
        new_data = np.array(
            [pd.to_numeric(flange_limits[column].values, errors="coerce") for column in flange_limits.columns]
        ).transpose()
        flange_limits = pd.DataFrame(columns=flange_limits.columns, data=new_data).fillna(0)

        # make sure pressures are positive
        if not all(flange_limits.loc[:, flange_limits.columns != "Temperature"].fillna(0).values.flatten() >= 0):
            raise ValueError("\nPressure less than zero.")

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            # add units to temperature column
            flange_limits["Temperature"] = [_Q(temp, "degC") for temp in flange_limits["Temperature"]]

            # add units to pressure columns
            for key in flange_limits.keys():
                if key != "Temperature":
                    pressures = []
                    for pressure in flange_limits[key]:
                        if pressure < 0:
                            pressures.append(np.NaN)
                        else:
                            pressures.append(_Q(float(pressure), "bar"))
                    flange_limits[key] = pressures

        group_flange_limits[group] = flange_limits
    return group_flange_limits

FLANGE_LIMITS = _get_flange_limits_from_csv()

def _import_thread_specs():
    """
    Imports thread specifications from .csv files

    Returns
    -------
    thread_specs : dict
        [internal thread specs, external thread specs]. Both sets of thread
        specifications are multi-indexed with (thread size, thread class).
    """
    file_directory = os.path.join(os.path.dirname(os.path.relpath(__file__)), "lookup_data")
    file_names = ["ANSI_inch_internal_thread.csv", "ANSI_inch_external_thread.csv"]
    file_locations = [os.path.relpath(os.path.join(file_directory, name)) for name in file_names]

    thread_specs = {
        key: pd.read_csv(location, index_col=(0, 1)) for location, key in zip(file_locations, ["internal", "external"])
    }

    return thread_specs

THREAD_SPECS = _import_thread_specs()

class Bolt:
    """
    Methods relating to bolt calculations and property lookup.
    """

    @classmethod
    def calculate_safety_factors(
        cls,
        max_pressure,
        window_area,
        num_bolts,
        thread_size,
        thread_class,
        bolt_max_tensile,
        plate_max_tensile,
        engagement_length,
        unit_registry=_U,
    ):
        """
        Calculates bolt and plate safety factors for viewing window bolts

        Parameters
        ----------
        max_pressure : pint.Quantity or Tuple[float, str]
            Tube maximum pressure
        window_area : pint.Quantity or Tuple[float, str]
            Window area exposed to high pressure environment
        num_bolts : int or np.ndarray
            Number of bolts used to secure each viewing window
        thread_size : str
            Size of threads to be evaluated, e.g. ``1/4-20`` or ``1 1/2-6``
        thread_class : str
            Class of threads to be evaluated, ``"2"`` or ``"3"``. (``"A"`` or
            ``"B"`` are automatically appended for internal/external threads)
        bolt_max_tensile : pint.Quantity or Tuple[float, str]
            Pint quantity of bolt (ext. thread) tensile failure stress
        plate_max_tensile : pint.Quantity or Tuple[float, str]
            Pint quantity of plate (int. thread) tensile failure stress
        engagement_length : pint.Quantity or Tuple[float, str]
            Pint quantity of total thread engagement length
        unit_registry : pint.UnitRegistry
            Keeps output consistent with parent registry, avoiding conflicts

        Returns
        -------
        dict
            Dictionary giving factors of safety for window bolts and the plate

            that they are screwed into. Keys:

            * ``"bolt"``
            * ``"plate"``
        """
        quant = unit_registry.Quantity

        max_pressure = units.parse_quant_input(max_pressure, unit_registry)
        window_area = units.parse_quant_input(window_area, unit_registry)
        bolt_max_tensile = units.parse_quant_input(bolt_max_tensile, unit_registry)
        plate_max_tensile = units.parse_quant_input(plate_max_tensile, unit_registry)
        engagement_length = units.parse_quant_input(engagement_length, unit_registry)

        units.check_pint_quantity(max_pressure, "pressure", ensure_positive=True)
        units.check_pint_quantity(window_area, "area", ensure_positive=True)
        units.check_pint_quantity(bolt_max_tensile, "pressure", ensure_positive=True)
        units.check_pint_quantity(plate_max_tensile, "pressure", ensure_positive=True)
        units.check_pint_quantity(engagement_length, "length", ensure_positive=True)

        # get total force per bolt
        window_force = (max_pressure - quant(1, "atm")) * window_area / num_bolts

        # get stress areas
        thread = cls.calculate_stress_areas(
            thread_size,
            thread_class,
            bolt_max_tensile,
            plate_max_tensile,
            engagement_length,
            unit_registry,
        )
        screw_area = thread["screw area"]
        screw_area = quant(screw_area.magnitude, screw_area.units.format_babel())
        plate_area = thread["plate area"]
        plate_area = quant(plate_area.magnitude, plate_area.units.format_babel())

        # calculate safety factors
        safety_factors = dict()
        safety_factors["bolt"] = (bolt_max_tensile / (window_force / screw_area)).to("").magnitude
        safety_factors["plate"] = (plate_max_tensile / (window_force / plate_area)).to("").magnitude
        return safety_factors

    @staticmethod
    def calculate_stress_areas(
        thread_size,
        thread_class,
        bolt_max_tensile,
        plate_max_tensile,
        engagement_length,
        unit_registry,
    ):
        """
        Calculates internal and external thread stress areas using formulas in
        Machinery's Handbook, 26th edition.

        Parameters
        ----------
        thread_size : str
            Size of threads to be evaluated, e.g. ``"1/4-20"`` or ``"1 1/2-6"``
        thread_class : str
            Class of threads to be evaluated, ``"2"`` or ``"3"``. (``"A"`` or
            ``"B"`` are automatically appended for internal/external threads)
        bolt_max_tensile : pint.Quantity or Tuple[float, str]
            Pint quantity of bolt (external thread) tensile failure stress
        plate_max_tensile : pint.Quantity or Tuple[float, str]
            Pint quantity of plate (internal thread) tensile failure stress
        engagement_length : pint.Quantity or Tuple[float, str]
            Pint quantity of total thread engagement length
        unit_registry : pint.UnitRegistry
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        dict
            Dictionary with the following keys:

            * ``"plate area"``: Stress area of internal threads within the plate
            * ``"screw area"``: Stress area of external threads on the screw
            * ``"minimum engagement"``: Minimum engagement length causing screw
              to fail in tension rather than shear, thus preventing the plate
              from stripping.
        """
        quant = unit_registry.Quantity

        units.check_pint_quantity(bolt_max_tensile, "pressure", ensure_positive=True)
        units.check_pint_quantity(plate_max_tensile, "pressure", ensure_positive=True)
        units.check_pint_quantity(engagement_length, "length", ensure_positive=True)

        # convert to local unit registry
        bolt_max_tensile = quant(bolt_max_tensile.magnitude, bolt_max_tensile.units.format_babel())
        plate_max_tensile = quant(plate_max_tensile.magnitude, plate_max_tensile.units.format_babel())
        engagement_length = quant(engagement_length.magnitude, engagement_length.units.format_babel())

        thread = dict()

        # look up thread specs for stress area calculations
        k_n_max = quant(
            THREAD_SPECS["internal"]["minor diameter max"][thread_size][thread_class + "B"],
            "in",
        )
        e_s_min = quant(
            THREAD_SPECS["external"]["pitch diameter min"][thread_size][thread_class + "A"],
            "in",
        )
        e_n_max = quant(
            THREAD_SPECS["internal"]["pitch diameter max"][thread_size][thread_class + "B"],
            "in",
        )
        d_s_min = quant(
            THREAD_SPECS["external"]["major diameter min"][thread_size][thread_class + "A"],
            "in",
        )
        tpi = quant(float(thread_size.split("-")[-1]), "1/in")
        basic_diameter = quant(
            THREAD_SPECS["external"]["basic diameter"][thread_size][thread_class + "A"],
            "in",
        )

        if bolt_max_tensile < quant(100000, "psi"):
            # calculate screw tensile area using eq. 9 (p. 1482) in Fasteners
            # section of Machinery's Handbook 26 (also eq. 2a on p. 1490)
            screw_area_tensile = np.pi / 4 * (basic_diameter - 0.9742785 / tpi) ** 2
        else:
            # calculate screw tensile area using eq. 2b (p. 1490) in Fasteners
            # section of Machinery's Handbook 26
            screw_area_tensile = np.pi * (e_s_min / 2 - 0.16238 / tpi) ** 2

        # calculate screw shear area using eq. 5 (p. 1491) in Fasteners section
        # of Machinery's Handbook 26
        screw_area_shear = np.pi * tpi * engagement_length * k_n_max * (1.0 / (2 * tpi) + 0.57735 * (e_s_min - k_n_max))

        # choose correct area
        if screw_area_shear < screw_area_tensile:
            warnings.warn(
                "Screws fail in shear, not tension."
                + " Plate may be damaged."
                + " Consider increasing bolt engagement length",
                Warning,
            )
            thread["screw area"] = screw_area_shear
        else:
            thread["screw area"] = screw_area_tensile

        # calculate plate shear area using eq. 6 (p. 1491) in Fasteners section
        # of Machinery's Handbook 26
        thread["plate area"] = (
            np.pi * tpi * engagement_length * d_s_min * (1.0 / (2 * tpi) + 0.57735 * (d_s_min - e_n_max))
        )

        # calculate minimum engagement scale factor using eq. 3 (p. 1490) in
        # Fasteners section of Machinery's Handbook 26
        j_factor = (screw_area_shear * bolt_max_tensile) / (thread["plate area"] * plate_max_tensile)

        # calculate minimum thread engagement (corrected for material
        # differences) using eqs. 1 and 4 (pp. 1490-1491) in Fasteners section
        # of Machinery's Handbook 26
        thread["minimum engagement"] = (
            2 * screw_area_tensile / (k_n_max * np.pi * (1.0 / 2 + 0.57735 * tpi * (e_s_min - k_n_max)))
        ) * j_factor

        return thread

    @staticmethod
    def get_thread_property(thread_property, thread_size, thread_class, unit_registry):
        """
        Finds a thread property, such as minor diameter, using a dataframe from
        import_thread_specs(). import_thread_specs is not directly called here
        to save time by not reading from disk every time a property is
        requested.

        Parameters
        ----------
        thread_property : str
            Property that is desired, such as ``"minor diameter"``
        thread_size : str
            Thread size for desired property, such as ``"1/4-20"`` or
            ``"1 1/2-6"``
        thread_class : str
            Thread class: ``"2B"`` or ``"3B"`` for internal threads, ``"2A"`` or
            ``"3A"`` for external threads
        unit_registry : pint.UnitRegistry
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        pint.Quantity
            Property requested, as a pint quantity with units of inches
        """
        quant = unit_registry.Quantity

        # determine if internal or external
        if "A" in thread_class and ("2" in thread_class or "3" in thread_class):
            thread_specs = THREAD_SPECS["external"]
        elif "B" in thread_class and ("2" in thread_class or "3" in thread_class):
            thread_specs = THREAD_SPECS["internal"]
        else:
            raise ValueError("\nbad thread class")

        # ensure property is in the specs dataframe
        if thread_property not in thread_specs.keys():
            raise KeyError(
                "\nThread property '"
                + thread_property
                + "' not found. Available specs: "
                + "'"
                + "', '".join(thread_specs.keys())
                + "'"
            )

        # ensure thread size is in the specs dataframe
        if thread_size not in thread_specs.index:
            raise KeyError("\nThread size '" + thread_size + "' not found")

        # retrieve the property
        return quant(thread_specs[thread_property][thread_size][thread_class], "in")

class DDT:
    """
    Methods for estimating the deflagration-to-detonation transition (DDT).
    """

    @staticmethod
    def calculate_blockage_diameter(pipe_id, blockage_ratio, unit_registry=_U):
        """
        Calculates the diameter of a Shchelkin spiral corresponding to a given
        blockage ratio within a pipe of given inner diameter.

        Parameters
        ----------
        pipe_id : pint.Quantity or Tuple[float, str]
            Length scale representing the inner diameter of the pipe used for
            the detonation tube
        blockage_ratio : float
            Ratio of blocked area to total cross-sectional area, :math:`0 < BR
            < 1`
        unit_registry : pint.UnitRegistry
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        pint.Quantity
            Shchelkin spiral diameter inside a tube of `pipe_id` inner diameter
            giving a blockage ratio of `blockage_ratio`. Units are the same as
            pipe_id.
        """
        # ensure blockage ratio is a float
        try:
            blockage_ratio = float(blockage_ratio)
        except ValueError:
            raise ValueError("\nNon-numeric blockage ratio.")

        # ensure blockage ratio is on 0<BR<1
        if not 0 < blockage_ratio < 1:
            raise ValueError("\nBlockage ratio outside of 0<BR<1")

        pipe_id = units.parse_quant_input(pipe_id, unit_registry)
        units.check_pint_quantity(pipe_id, "length", ensure_positive=True)

        # calculate Shchelkin spiral diameter
        spiral_diameter = pipe_id / 2 * (1 - sqrt(1 - blockage_ratio))
        return spiral_diameter

    @staticmethod
    def calculate_blockage_ratio(tube_inner_diameter, blockage_diameter, unit_registry=_U):
        """
        Calculates the blockage ratio of a Shchelkin spiral within a detonation
        tube.

        Parameters
        ----------
        tube_inner_diameter : pint.Quantity or Tuple[float, str]
            Inner diameter of the detonation tube
        blockage_diameter : pint.Quantity or Tuple[float, str]
            Outer diameter of the blockage used to create the Shchelkin spiral
            (i.e. a Shchelkin spiral made from 1/2" round stock would be
            ``blockage_diameter=quant(0.5, "inch")``
        unit_registry : pint.UnitRegistry
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        float
            Ratio of blocked to open area (between 0 and 1)
        """

        tube_inner_diameter = units.parse_quant_input(tube_inner_diameter, unit_registry).to_base_units()
        blockage_diameter = units.parse_quant_input(blockage_diameter, unit_registry).to_base_units()
        units.check_pint_quantity(tube_inner_diameter, "length", ensure_positive=True)
        units.check_pint_quantity(blockage_diameter, "length", ensure_positive=True)

        # ensure blockage diameter < tube diameter
        if tube_inner_diameter.magnitude == 0:
            raise ValueError("\ntube ID cannot be 0")
        elif blockage_diameter >= tube_inner_diameter:
            raise ValueError("\nblockage diameter >= tube diameter")

        # calculate blockage ratio
        blockage_ratio = 1 - (1 - 2 * blockage_diameter.magnitude / tube_inner_diameter.magnitude) ** 2

        return blockage_ratio

    @staticmethod
    def calculate_run_up(
        blockage_ratio,
        tube_diameter,
        initial_temperature,
        initial_pressure,
        species_dict,
        mechanism,
        unit_registry=_U,
        phase_specification="",
    ):
        """
        Calculates the runup distance needed for a detonation to develop from a
        deflagration for a given blockage ratio, tube diameter, and mixture.
        This is accomplished using equations collected by Ciccarelli and
        Dorofeev [1] for blockage ratios <= 0.75. If the desired blockage ratio
        is less than 0.3, the mixture viscosity is needed, and the
        `phase_specification` option may be necessary depending on the
        mechanism.

        [1] G. Ciccarelli and S. Dorofeev, *Flame acceleration and transition to
        detonation in ducts*, Progress in Energy and Combustion Science,
        vol. 34, no. 4, pp. 499–550, Aug. 2008.
        https://doi.org/10.1016/j.pecs.2007.11.002

        Parameters
        ----------
        blockage_ratio : float
            Ratio of blocked area to total cross-sectional area, :math:`0 < BR
            < 1`
        tube_diameter : pint.Quantity or Tuple[float, str]
            Internal diameter of the detonation tube
        initial_temperature : pint.Quantity or Tuple[float, str]
            Mixture initial temperature
        initial_pressure : pint.Quantity or Tuple[float, str]
            Mixture initial pressure
        species_dict : dict
            Dictionary containing the species in the mixture as keys, with total
            moles or mole fractions as values
        mechanism : str
            Mechanism file name for Cantera. See ``thermochem.find_mechanisms()`` for
            a list of installed mechanisms.
        unit_registry : pint.UnitRegistry
            Unit registry for managing units to prevent conflicts with parent
            unit registry
        phase_specification : str, optional
            Phase specification within the mechanism file used to evaluate
            thermophysical properties. If ``gri30.yaml`` is used with no phase
            specification viscosity calculations will fail, resulting in
            an error for all blockage ratios less than 0.3.

        Returns
        -------
        pint.Quantity
            Predicted DDT distance
        """

        if blockage_ratio <= 0 or blockage_ratio > 0.75:
            raise ValueError("\nBlockage ratio outside of correlation range")

        tube_diameter = units.parse_quant_input(tube_diameter, unit_registry)
        initial_temperature = units.parse_quant_input(initial_temperature, unit_registry)
        initial_pressure = units.parse_quant_input(initial_pressure, unit_registry)
        units.check_pint_quantity(tube_diameter, "length", ensure_positive=True)
        units.check_pint_quantity(initial_temperature, "temperature", ensure_positive=True)
        units.check_pint_quantity(initial_pressure, "pressure", ensure_positive=True)

        # calculate laminar flame speed
        laminar_fs = thermochem.calculate_laminar_flame_speed(
            initial_temperature,
            initial_pressure,
            species_dict,
            mechanism,
            phase_specification,
            unit_registry,
        )

        # calculate density ratio across the deflagration assuming adiabatic
        # flame
        density = np.zeros(2)
        working_gas = ct.Solution(mechanism, phase_specification)
        working_gas.TPX = [
            initial_temperature.to("K").magnitude,
            initial_pressure.to("Pa").magnitude,
            species_dict,
        ]
        density[0] = working_gas.density
        working_gas.equilibrate("HP")
        density[1] = 1 / working_gas.density
        density_ratio = np.prod(density)

        # find sound speed in products at adiabatic flame temperature
        sound_speed = thermochem.get_eq_sound_speed(
            (working_gas.T, "K"),
            (working_gas.P, "Pa"),
            species_dict,
            mechanism,
            phase_specification,
            unit_registry,
        )

        def eq4_1():
            """
            Calculate runup distance for blockage ratios <= 0.1 using equation
            4.1 from G. Ciccarelli and S. Dorofeev, “Flame acceleration and
            transition to detonation in ducts,” Progress in Energy and
            Combustion Science, vol. 34, no. 4, pp. 499–550, Aug. 2008.
            """
            # define constants
            kappa = 0.4
            kk = 5.5
            cc = 0.2
            mm = -0.18
            eta = 2.1

            # calculate laminar flame thickness, delta
            working_gas.TPX = [
                initial_temperature.to("K").magnitude,
                initial_pressure.to("Pa").magnitude,
                species_dict,
            ]
            rho = units.parse_quant_input((working_gas.density_mass, "kg/m^3"), unit_registry)
            mu = units.parse_quant_input((working_gas.viscosity, "Pa*s"), unit_registry)
            nu = mu / rho
            delta = (nu / laminar_fs).to_base_units()

            # calculate gamma
            gamma = (
                sound_speed / (eta * (density_ratio - 1) ** 2 * laminar_fs) * (delta / tube_diameter) ** (1.0 / 3)
            ) ** (1 / (2 * mm + 7.0 / 3))

            # calculate runup distance
            d_over_h = 2.0 / (1 - np.sqrt(1 - blockage_ratio))
            runup = gamma / cc * (1 / kappa * np.log(gamma * d_over_h) + kk) * tube_diameter
            return runup.to(tube_diameter.units.format_babel())

        def eq4_4():
            """
            Calculate runup for blockage ratios between 0.3 and 0.75 using
            equation 4.4 in G. Ciccarelli and S. Dorofeev, “Flame acceleration
            and transition to detonation in ducts,” Progress in Energy and
            Combustion Science, vol. 34, no. 4, pp. 499–550, Aug. 2008.
            """
            # define constants
            aa = 2.0
            bb = 1.5

            # calculate left and right hand sides of eq 4.4
            lhs = 2 * 10 * laminar_fs * (density_ratio - 1) / (sound_speed * tube_diameter)
            rhs = aa * (1 - blockage_ratio) / (1 + bb * blockage_ratio)

            runup = rhs / lhs

            return runup.to(tube_diameter.units.format_babel())

        # use appropriate equation to calculate runup distance
        if 0.3 <= blockage_ratio <= 0.75:
            runup_distance = eq4_4()
        elif 0.1 >= blockage_ratio:
            runup_distance = eq4_1()
        else:
            interp_distances = np.array([eq4_1().magnitude, eq4_4().magnitude])
            runup_distance = np.interp(blockage_ratio, np.array([0.1, 0.3]), interp_distances)
            runup_distance = units.parse_quant_input(
                (runup_distance, tube_diameter.units.format_babel()), unit_registry
            )

        return runup_distance

class Window:
    """
    Methods for the design of viewing windows for optical access.
    """

    @classmethod
    def safety_factor(cls, length, width, thickness, pressure, rupture_modulus, unit_registry=_U):
        """
        Calculates the safety factor of a clamped rectangular
        window given window dimensions, design pressure, and material rupture
        modulus

        Parameters
        ----------
        length : pint.Quantity or Tuple[float, str]
            Window unsupported (viewing) length
        width : pint.Quantity or Tuple[float, str]
            Window unsupported (viewing) width
        thickness : pint.Quantity or Tuple[float, str]
            Window thickness
        pressure : pint.Quantity or Tuple[float, str]
            Design pressure differential across window at which factor of
            safety is to be calculated
        rupture_modulus : pint.Quantity or Tuple[float, str]
            Rupture modulus of desired window material
        unit_registry : pint.UnitRegistry
            Keeps output consistent with parent registry, avoiding conflicts

        Returns
        -------
        float
            Window factor of safety
        """
        length = units.parse_quant_input(length, unit_registry).to_base_units()
        width = units.parse_quant_input(width, unit_registry).to_base_units()
        thickness = units.parse_quant_input(thickness, unit_registry).to_base_units()
        pressure = units.parse_quant_input(pressure, unit_registry).to_base_units()
        rupture_modulus = units.parse_quant_input(rupture_modulus, unit_registry).to_base_units()

        units.check_pint_quantity(length, "length", ensure_positive=True)
        units.check_pint_quantity(width, "length", ensure_positive=True)
        units.check_pint_quantity(thickness, "length", ensure_positive=True)

        units.check_pint_quantity(pressure, "pressure", ensure_positive=True)
        units.check_pint_quantity(rupture_modulus, "pressure", ensure_positive=True)

        # https://www.crystran.co.uk/userfiles/files/design-of-pressure-windows.pdf
        k = 0.75
        r = length / width
        safety_factor = (
            ((thickness / length) ** 2 * (2 * rupture_modulus * (1 + r**2)) / (k * pressure)).to("").magnitude
        )
        return safety_factor

    @classmethod
    def minimum_thickness(cls, length, width, safety_factor, pressure, rupture_modulus, unit_registry=_U):
        """
        Calculates the thickness of a clamped rectangular window which gives
        the desired safety factor

        Parameters
        ----------
        length : pint.Quantity or Tuple[float, str]
            Window unsupported (viewing) length
        width : pint.Quantity or Tuple[float, str]
            Window unsupported (viewing) width
        safety_factor : float
            Safety factor
        pressure : pint.Quantity or Tuple[float, str]
            Design pressure differential across window at which factor of
            safety is to be calculated
        rupture_modulus : pint.Quantity or Tuple[float, str]
            Rupture modulus of desired window material
        unit_registry : pint.UnitRegistry
            Keeps output consistent with parent registry, avoiding conflicts

        Returns
        -------
        pint.Quantity
            Window thickness
        """
        length = units.parse_quant_input(length, unit_registry).to_base_units()
        width = units.parse_quant_input(width, unit_registry).to_base_units()
        pressure = units.parse_quant_input(pressure, unit_registry).to_base_units()
        rupture_modulus = units.parse_quant_input(rupture_modulus, unit_registry).to_base_units()

        units.check_pint_quantity(length, "length", ensure_positive=True)
        units.check_pint_quantity(width, "length", ensure_positive=True)
        units.check_pint_quantity(pressure, "pressure", ensure_positive=True)
        units.check_pint_quantity(rupture_modulus, "pressure", ensure_positive=True)

        # Ensure safety factor is numeric and > 1
        try:
            if safety_factor < 1:
                raise ValueError("\nWindow safety factor < 1")
        except TypeError:
            raise TypeError("\nNon-numeric window safety factor")

        # https://www.crystran.co.uk/userfiles/files/design-of-pressure-windows.pdf
        k = 0.75  # clamped
        r = length / width
        thickness = (length * np.sqrt(safety_factor * k / 2) * np.sqrt(pressure / (rupture_modulus * (1 + r**2)))).to(
            width.units.format_babel()
        )

        return thickness

class Tube:
    """
    Methods for designing and determining operational limits of a detonation
    tube.

    For a list of available materials see ``Tube.available_materials``
    """

    available_pipe_sizes = list(PIPE_SCHEDULES.index.values)
    available_materials = list(TUBE_MATERIALS.Grade.values)

    @classmethod
    def calculate_max_stress(cls, initial_temperature, material, welded, unit_registry=thermochem._U):
        """
        Finds the maximum allowable stress of a tube material at the tube's
        initial temperature

        Parameters
        ----------
        initial_temperature : pint.Quantity or Tuple[float, str]
        material : str
            Pipe material (see ``Tube.available_materials``)
        welded : bool
            ``True`` for welded pipe; ``False`` for seamless
        unit_registry : pint.UnitRegistry
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        pint.Quantity
            Maximum allowable tube stress
        """
        initial_temperature = units.parse_quant_input(initial_temperature, unit_registry)

        # look up stress-temperature limits and units
        stress_limits = cls._get_pipe_stress_limits(material=material, welded=welded, unit_registry=unit_registry)
        stress_units = stress_limits.values[0].units
        stresses = list(map(lambda x: x.magnitude, stress_limits.values))
        temp_units = stress_limits.index[0].units
        temperatures = list(map(lambda x: x.magnitude, stress_limits.index.values))

        # ensure material stress limits have monotonically increasing
        # temperatures, otherwise the np.interp "results are nonsense" per
        # scipy docs
        if not np.all(np.diff(temperatures) > 0):
            raise ValueError("\nStress limits require temperatures to be " + "monotonically increasing")

        # interpolate max stress
        max_stress = unit_registry.Quantity(
            np.interp(initial_temperature.to(temp_units).magnitude, temperatures, stresses),
            stress_units,
        )

        return max_stress

    @staticmethod
    def _get_pipe_stress_limits(material, welded, unit_registry=_U):
        """
        Looks up ASME B31.1 stress limits as a function of pipe temperature

        Parameters
        ----------
        material : str
            Pipe material (see ``Tube.available_materials``)
        welded : bool
            ``True`` for welded pipe; ``False`` for seamless
        unit_registry : pint.UnitRegistry, optional
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        pd.Series
            Pandas series with data consisting of tress limits (as pint
            quantities), and indices of the corresponding temperatures (also
            as pint quantities)
        """
        quant = unit_registry.Quantity
        if welded:
            material_limits = MATERIAL_LIMITS["welded"][["Temp", material]]
        else:
            material_limits = MATERIAL_LIMITS["seamless"][["Temp", material]]
        material_limits.columns = ["temperature", "stress"]

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            material_limits = pd.Series(
                data=[quant(t, "ksi") for t in material_limits["stress"].values],
                index=[quant(t, "degF") for t in material_limits["temperature"].values],
                name="stress",
            )

        return material_limits

    @staticmethod
    def dynamic_load_factor(
        tube_id,
        tube_od,
        cj_velocity,
        elastic_modulus,
        density,
        poisson_ratio,
        plus_or_minus=0.1,
    ):
        """
        Calculates the dynamic load factor (DLF) by which a detonation tube's
        static analysis should be scaled in order to account for the tube's
        response to transient pressures. DLF calculation is based on the work of
        Shepherd [1]. Since the limits of "approximately equal to" are not
        defined in the paper, a default value of plus or minus ten percent
        is assumed, thus `plus_or_minus=0.1`.

        [1] Shepherd, J. E. (2009). *Structural Response of Piping to
        Internal Gas Detonation*. Journal of Pressure Vessel Technology,
        131(3), 031204. https://doi.org/10.1115/1.3089497

        Parameters
        ----------
        tube_id : pint.Quantity or Tuple[float, str]
            Inner diameter of tube
        tube_od : pint.Quantity or Tuple[float, str]
            Outer diameter of tube
        cj_velocity : pint.Quantity or Tuple[float, str]
            CJ velocity
        elastic_modulus : pint.Quantity or Tuple[float, str]
            Elastic modulus of tube material
        density : pint.Quantity or Tuple[float, str]
            Density of tube material
        poisson_ratio : float
            Poisson ratio of tube material
        plus_or_minus : float, optional
            Defines the band about the critical velocity which is considered
            "approximately equal to" -- the default value of 0.1 means plus
            or minus ten percent

        Returns
        -------
        float
            Factor by which the tube's static maximum pressure should be
            de-rated to account for transient response to detonation waves
        """
        units.check_pint_quantity(tube_id, "length", True)
        units.check_pint_quantity(tube_od, "length", True)
        units.check_pint_quantity(cj_velocity, "velocity")
        units.check_pint_quantity(elastic_modulus, "pressure", True)
        units.check_pint_quantity(density, "density", True)

        if not (0 < plus_or_minus < 1):
            raise ValueError("\nplus_or_minus factor not between 0 and 1")

        # set geometry
        # first /2 for averaging
        # second /2 to to convert diameter to radius
        radius = (tube_od + tube_id) / 4.0
        tube_thickness = (tube_od - tube_id) / 2.0

        # calculate critical velocity
        crit_velocity = (
            (elastic_modulus**2 * tube_thickness**2)
            / (3.0 * density**2 * radius**2 * (1.0 - poisson_ratio**2))
        ) ** (1.0 / 4)

        # set limits for "approximately Vcrit"
        bounds = [
            (crit_velocity * (1.0 + plus_or_minus)).to("m/s").magnitude,
            (crit_velocity * (1.0 - plus_or_minus)).to("m/s").magnitude,
        ]
        cj_velocity = cj_velocity.to("m/s").magnitude

        if cj_velocity < bounds[1]:
            dynamic_load_factor = 1
        elif cj_velocity > bounds[0]:
            dynamic_load_factor = 2
        else:
            dynamic_load_factor = 4

        return dynamic_load_factor

    @staticmethod
    def calculate_max_pressure(tube_id, tube_od, max_stress, safety_factor=1):
        """
        Calculates the maximum allowable pressure from the tube dimensions
        and stress limits using the basic longitudinal joint formula [1].

        [1] E. F. Megyesy, *Pressure vessel handbook*, Oklahoma City, OK:
        PV Publishing, Inc., 2001, p. 14.

        Parameters
        ----------
        tube_id : pint.Quantity or Tuple[float, str]
            Inner diameter of tube
        tube_od : pint.Quantity or Tuple[float, str]
            Outer diameter of tube
        max_stress : pint.Quantity or Tuple[float, str]
            Maximum allowable stress in tube material (e.g. from ASME B31.1)
        safety_factor : float, optional
            Desired safety factor

        Returns
        -------
        pint.Quantity
            Pressure resulting in maximum allowable stress
        """
        mean_diameter = (tube_od + tube_id) / 2.0
        wall_thickness = (tube_od - tube_id) / 2.0
        max_pressure = max_stress * (2 * wall_thickness) / (mean_diameter * safety_factor)

        return max_pressure

    @classmethod
    def calculate_max_initial_pressure(
        cls,
        tube_id,
        tube_od,
        initial_temperature,
        species_dict,
        mechanism,
        max_pressure,
        elastic_modulus,
        density,
        poisson_ratio,
        plus_or_minus=0.1,
        use_dynamic_load_factor=True,
        unit_registry=pint.UnitRegistry(),
        error_tol=1e-4,
        max_iterations=500,
        use_multiprocessing=False,
        max_pressure_state="reflected",
        return_dlf=False,
    ):
        """
        Parameters
        ----------
        tube_id : pint.Quantity or Tuple[float, str]
            Inner diameter of tube
        tube_od : pint.Quantity or Tuple[float, str]
            Outer diameter of tube
            Internal diameter of the detonation tube
        initial_temperature : pint.Quantity or Tuple[float, str]
            Mixture initial temperature
        species_dict : dict
            Dictionary containing the species in the mixture as keys, with total
            moles or mole fractions as values
        mechanism : str
            Mechanism file name for Cantera
        max_pressure : pint.Quantity or Tuple[float, str]
            Maximum allowable pressure within the tube
        elastic_modulus : pint.Quantity or Tuple[float, str]
            Elastic modulus of tube material
        density : pint.Quantity or Tuple[float, str]
            Density of tube material
        poisson_ratio : float
            Poisson ratio of tube material
        plus_or_minus : float
            For dynamic load factor calculation, ``plus_or_minus`` defines
            the band about the critical velocity which is considered
            "approximately equal to." The default value of 0.1 means plus
            or minus ten percent.
        use_dynamic_load_factor : bool
            ``True`` to use dynamic load factor, ``False`` to ignore it.
        unit_registry : pint unit registry
            Unit registry for managing units to prevent conflicts with parent
            unit registry
        error_tol : float
            Relative error tolerance below which initial pressure calculations
            are considered 'good enough'
        max_iterations : int
            Maximum number of loop iterations before exit, defaults to 500
        use_multiprocessing : bool
            ``True`` to use multiprocessing. Defaults to ``False``.
        max_pressure_state : str
            Which state to use as the tube's maximum pressure. Must be either
            ``"reflected"`` (closed tube) or ``"cj"`` (open tube).
        return_dlf : bool
            If ``True``, dynamic load factor will be returned

        Returns
        -------
        pint.Quantity or Tuple[pint.Quantity, float]
            Initial mixture pressure corresponding to the tube's maximum
            allowable pressure.

            If `return_dlf` is set to ``True``, this will be a list, with item
            0 being the initial pressure and item 1 being DLF
        """
        if max_pressure_state not in ("reflected", "cj"):
            raise ValueError("`max_pressure_state` must be 'reflected' or 'cj'")

        quant = unit_registry.Quantity

        # get a rough estimate of the initial pressure
        initial_pressure = quant(1, "atm")
        counter = 0
        error_tol = abs(error_tol)

        state = thermochem.calculate_reflected_shock_state(
            initial_temperature,
            initial_pressure,
            species_dict,
            mechanism,
            unit_registry,
            use_multiprocessing,
        )

        dlf = cls.dynamic_load_factor(
            tube_id,
            tube_od,
            state["cj"]["speed"],
            elastic_modulus,
            density,
            poisson_ratio,
            plus_or_minus,
        )

        error = (initial_pressure.magnitude * dlf / max_pressure.magnitude) - 1.0
        while abs(error) > error_tol and counter < max_iterations:
            counter += 1

            # update initial pressure guess
            initial_pressure = initial_pressure * max_pressure.magnitude / (dlf * state["reflected"]["state"].P)

            # get reflected shock pressure
            state = thermochem.calculate_reflected_shock_state(
                initial_temperature,
                initial_pressure,
                species_dict,
                mechanism,
                unit_registry,
                use_multiprocessing,
            )

            if use_dynamic_load_factor:
                dlf = cls.dynamic_load_factor(
                    tube_id,
                    tube_od,
                    state["cj"]["speed"],
                    elastic_modulus,
                    density,
                    poisson_ratio,
                    plus_or_minus,
                )
            else:
                dlf = 1.0

            error = (state[max_pressure_state]["state"].P * dlf - max_pressure.magnitude) / max_pressure.magnitude

        if return_dlf:
            return initial_pressure, dlf
        else:
            return initial_pressure

    @staticmethod
    def get_available_pipe_schedules(pipe_size):
        """
        Gets available pipe schedules for a given nominal size

        Parameters
        ----------
        pipe_size : str
            Nominal pipe size (see ``Tube.available_pipe_sizes``)

        Returns
        -------
        list
            List of available pipe schedules
        """
        if pipe_size not in PIPE_SCHEDULES.index:
            msg = "Invalid pipe size: %s. " "See Tube.available_pipe_sizes." % str(pipe_size)
            raise ValueError(msg)
        else:
            return list(PIPE_SCHEDULES.T[pipe_size].dropna().index.drop("OD"))

    @classmethod
    def get_dimensions(cls, pipe_size, pipe_schedule, unit_registry=_U):
        """

        Parameters
        ----------
        pipe_size : str
            Nominal pipe size (see ``Tube.available_pipe_sizes``)
        pipe_schedule : str
            Pipe schedule (see ``Tube.get_available_pipe_schedules`` for a list
            of available schedules for `pipe_size`)
        unit_registry : pint.UnitRegistry, optional
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        dict
            Dictionary of pipe dimensions with the keys:

            * ``"inner_diameter"``
            * ``"outer_diameter"``
            * ``"wall_thickness"``
        """
        # note: this also checks for valid size
        if pipe_schedule not in cls.get_available_pipe_schedules(pipe_size):
            msg = (
                "Schedule {:s} invalid for pipe size {:s}. See "
                "Tube.available_pipe_sizes and "
                "Tube.get_available_pipe_schedules.".format(str(pipe_schedule), str(pipe_size))
            )
            raise ValueError(msg)
        outer_diameter = PIPE_SCHEDULES["OD"][pipe_size]
        wall_thickness = PIPE_SCHEDULES[pipe_schedule][pipe_size]
        inner_diameter = outer_diameter - 2 * wall_thickness

        return dict(
            inner_diameter=unit_registry.Quantity(inner_diameter, "in"),
            outer_diameter=unit_registry.Quantity(outer_diameter, "in"),
            wall_thickness=unit_registry.Quantity(wall_thickness, "in"),
        )

    @classmethod
    def get_elastic_modulus(cls, material, unit_registry=_U):
        """
        Gets the elastic modulus for a valid tube material. For valid materials
        see ``Tube.available_materials``.

        Parameters
        ----------
        material : str
            Pipe material (see ``Tube.available_materials``)
        unit_registry : pint.UnitRegistry, optional
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        pint.Quantity
            Elastic modulus
        """
        cls._check_material(material)
        return units.parse_quant_input(
            TUBE_MATERIALS[TUBE_MATERIALS["Grade"] == material]["ElasticModulus"].values[0],
            unit_registry,
        )

    @classmethod
    def get_density(cls, material, unit_registry=_U):
        """
        Gets the density for a valid tube material. For valid materials see
        ``Tube.available_materials``.

        Parameters
        ----------
        material : str
            Pipe material (see ``Tube.available_materials``)
        unit_registry : pint.UnitRegistry, optional
            Unit registry for managing units to prevent conflicts with parent
            unit registry

        Returns
        -------
        pint.Quantity
            Density
        """
        cls._check_material(material)
        return units.parse_quant_input(
            TUBE_MATERIALS[TUBE_MATERIALS["Grade"] == material]["Density"].values[0],
            unit_registry,
        )

    @classmethod
    def get_poisson(cls, material):
        """
        Gets the Poisson ratio for a valid tube material. For valid materials
        see ``Tube.available_materials``.

        Parameters
        ----------
        material : str
            Pipe material (see ``Tube.available_materials``)

        Returns
        -------
        float
            Poisson ratio (unitless)
        """
        cls._check_material(material)
        return TUBE_MATERIALS[TUBE_MATERIALS["Grade"] == material]["Poisson"].values[0]

    @classmethod
    def _check_material(cls, material):
        """
        Checks a material to make sure it is in ``Tube.available_materials``.

        Parameters
        ----------
        material

        Returns
        -------
        """
        material = str(material)
        if material not in cls.available_materials:
            raise ValueError("Material %s not found" % material)

class Flange:
    """
    Methods pertaining to flange class and max pressure calculations

    For a list of available materials see ``Flange.available_materials``
    """

    available_materials = list(TUBE_MATERIALS.Grade[pd.notna(TUBE_MATERIALS.Group)].values)

    @staticmethod
    def _check_flange_class(flange_class, group):
        if flange_class not in FLANGE_LIMITS[group].keys():
            msg = "{:s} not found in group {:s}.".format(
                str(flange_class),
                str(group),
            )
            raise ValueError(msg)
        return True

    @staticmethod
    def get_class(max_pressure, temperature, material, unit_registry=_U):
        """
        Finds the minimum allowable flange class per ASME B16.5 for a given
        flange temperature and tube pressure.

        Parameters
        ----------
        max_pressure : pint.Quantity or Tuple[float, str]
            Maximum pressure within the system as a quantity or tuple of
            ``(magnitude, "units")``
        temperature : pint.Quantity or Tuple[float, str]
            Flange temperature as a quantity or tuple of
            ``(magnitude, "units")``
        material : str
            Flange material (see ``Flange.available_materials``)
        unit_registry : pint.UnitRegistry, optional
            Pint unit registry, if output within a particular registry is
            desired

        Returns
        -------
        str
            Minimum allowable flange class
        """
        max_pressure = units.parse_quant_input(max_pressure, unit_registry)
        temperature = units.parse_quant_input(temperature, unit_registry)
        units.check_pint_quantity(max_pressure, "pressure", True)
        units.check_pint_quantity(temperature, "temperature", True)
        max_pressure = units.parse_quant_input(max_pressure, _U)
        temperature = units.parse_quant_input(temperature, _U)

        _check_material(material)
        group = str(TUBE_MATERIALS[TUBE_MATERIALS.Grade == material].Group.values[0])
        df_limits = FLANGE_LIMITS[group]

        # locate max pressure
        class_keys = df_limits.keys()[1:]
        max_key = "0"
        for key in class_keys:
            if int(key) > int(max_key):
                max_key = key
        max_ok_pressure = df_limits[max_key].dropna().max()

        # ensure pressure is within bounds
        if max_pressure > max_ok_pressure:
            raise ValueError("Pressure out of range.")

        # ensure temperature is within bounds
        if (temperature < df_limits["Temperature"].min()) or (temperature > df_limits["Temperature"].max()):
            raise ValueError("Temperature out of range.")

        df_limits = df_limits.applymap(lambda x: x.to_base_units().magnitude)
        df_interp = pd.DataFrame(columns=df_limits.columns, dtype=np.float64)
        df_interp["Temperature"] = [temperature.to_base_units().magnitude]
        df_interp = pd.concat((df_limits, df_interp)).set_index("Temperature")
        with warnings.catch_warnings():
            # ignore max iteration warning during interpolation
            warnings.simplefilter("ignore")
            df_interp.interpolate(method="slinear", inplace=True)
        ser_result = df_interp.iloc[-1]
        return ser_result.keys()[ser_result >= max_pressure.to_base_units().magnitude][0]

    @staticmethod
    def get_max_pressure(flange_class, temperature, material):
        """
        Finds the max allowable pressure for a flange of a given material and
        class at the desired temperature.

        Parameters
        ----------
        flange_class : str
            Class of the flange to evaluate, e.g. ``"900"``
        temperature : pint.Quantity or Tuple[float, str]
            Flange temperature
        material : str
            Flange material (see ``Flange.available_materials``)

        Returns
        -------
        pint.Quantity
            Maximum allowable pressure
        """
        _check_material(material)
        group = str(TUBE_MATERIALS.Group[TUBE_MATERIALS.Grade == material].iloc[0])
        Flange._check_flange_class(flange_class, group)

        # ensure correct unit registry is used for comparison
        if hasattr(temperature, "_REGISTRY"):
            ureg_out = temperature._REGISTRY
        else:
            ureg_out = _U

        temperature = units.parse_quant_input(temperature, _U)

        df_limits = FLANGE_LIMITS[group]
        # ensure temperature is in range
        if (temperature < df_limits["Temperature"].min()) or (temperature > df_limits["Temperature"].max()):
            raise ValueError("Temperature out of range.")

        df_limits = df_limits.applymap(lambda x: x.to_base_units().magnitude)

        df_interp = pd.DataFrame(columns=df_limits.columns, dtype=np.float64)
        df_interp["Temperature"] = [temperature.to_base_units().magnitude]
        df_interp = pd.concat((df_limits, df_interp)).set_index("Temperature")
        with warnings.catch_warnings():
            # ignore max iteration warning during interpolation
            warnings.simplefilter("ignore")
            df_interp.interpolate(method="slinear", inplace=True)
        ser_result = df_interp.iloc[-1]
        return ureg_out.Quantity(ser_result[flange_class], "Pa")

    @staticmethod
    def get_max_temperature(flange_class, pressure, material):
        """
        Finds the max allowable temperature for a flange of a given material and
        class at the desired pressure.

        Parameters
        ----------
        flange_class : str
            Class of the flange to evaluate, e.g. ``"900"``
        pressure : pint.Quantity or Tuple[float, str]
            System pressure
        material : str
            Flange material (see ``Flange.available_materials``)

        Returns
        -------
        pint.Quantity
            Max allowable flange temperature
        """
        _check_material(material)
        group = str(TUBE_MATERIALS.Group[TUBE_MATERIALS.Grade == material].iloc[0])
        Flange._check_flange_class(flange_class, group)

        # ensure correct unit registry is used for comparison
        if hasattr(pressure, "_REGISTRY"):
            ureg_out = pressure._REGISTRY
        else:
            ureg_out = _U

        pressure = units.parse_quant_input(pressure, _U)

        df_limits = FLANGE_LIMITS[group]
        # ensure pressure is in range
        if (pressure < df_limits[flange_class].min()) or (pressure > df_limits[flange_class].max()):
            raise ValueError("Pressure out of range.")

        df_limits = df_limits.applymap(lambda x: x.to_base_units().magnitude)[["Temperature", flange_class]]

        df_interp = pd.DataFrame(columns=df_limits.columns, dtype=np.float64)
        df_interp[flange_class] = [pressure.to_base_units().magnitude]
        df_interp = pd.concat((df_limits, df_interp)).set_index(flange_class)
        with warnings.catch_warnings():
            # ignore max iteration warning during interpolation
            warnings.simplefilter("ignore")
            df_interp.interpolate(method="slinear", inplace=True)
        ser_result = df_interp.iloc[-1]
        return ureg_out.Quantity(ser_result["Temperature"], "K")


from datetime import datetime
import data_provider
import joblib
import mdm_model
import numpy as np
import os.path
import slim
import tensorflow as tf
import time
import utils
import menpo
import menpo.io as mio

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_float('initial_learning_rate', 0.001,
                          """Initial learning rate.""")
tf.app.flags.DEFINE_float('num_epochs_per_decay', 5.0,
                          """Epochs after which learning rate decays.""")
tf.app.flags.DEFINE_float('learning_rate_decay_factor', 0.1,
                          """Learning rate decay factor.""")
tf.app.flags.DEFINE_integer('batch_size', 80, """The batch size to use.""")
tf.app.flags.DEFINE_integer('num_preprocess_threads', 4,
                            """How many preprocess threads to use.""")
tf.app.flags.DEFINE_string('train_dir', '/root/sharedfolder/flm/mdm/mdm/ckpt/train_22_concat3_last256',
                           """Directory where to write event logs """
                           """and checkpoint.""")
tf.app.flags.DEFINE_string('pretrained_model_checkpoint_path', '/root/sharedfolder/flm/mdm/mdm/ckpt/train_22_concat3_last256/model.ckpt',
                           """If specified, restore this pretrained model """
                           """before beginning any training.""")
tf.app.flags.DEFINE_string('pretrained_model_step', 2100,
                           """If specified, restore this pretrained model step """
                           """before beginning any training.""")
tf.app.flags.DEFINE_integer('max_steps', 100000,
                            """Number of batches to run.""")
tf.app.flags.DEFINE_string('train_device', '/gpu:1', """Device to train with.""")
tf.app.flags.DEFINE_string('datasets', ':'.join(
    ('/root/sharedfolder/flm/datasets/lfpw/trainset/*.png',  #image_080
     '/root/sharedfolder/flm/datasets/helen/trainset/*.jpg',
     '/root/sharedfolder/flm/datasets/afw/*.jpg')),     #7003
                           """Directory where to write event logs """
                           """and checkpoint.""")
tf.app.flags.DEFINE_integer('patch_size', 22, 'The extracted patch size')
# The decay to use for the moving average.
MOVING_AVERAGE_DECAY = 0.9999

def train(scope=''):
    """Train on dataset for a number of steps."""
    with tf.Graph().as_default(), tf.device(FLAGS.train_device):
        # Create a variable to count the number of train() calls. This equals the
        # number of batches processed * FLAGS.num_gpus.
        global_step = tf.get_variable('global_step', [],
                                      initializer=tf.constant_initializer(0),
                                      trainable=False)

        train_dirs = FLAGS.datasets.split(':')

        # Calculate the learning rate schedule.
        decay_steps = 15000

        # Decay the learning rate exponentially based on the number of steps.
        lr = tf.train.exponential_decay(FLAGS.initial_learning_rate,
                                        global_step,
                                        decay_steps,
                                        FLAGS.learning_rate_decay_factor,
                                        staircase=True)

        # Create an optimizer that performs gradient descent.
        opt = tf.train.AdamOptimizer(lr)

        # Override the number of preprocessing threads to account for the increased
        # number of GPU towers.
        num_preprocess_threads = FLAGS.num_preprocess_threads

        _images, _shapes, _reference_shape, pca_model = \
            data_provider.load_images(train_dirs)

        reference_shape = tf.constant(_reference_shape,
                                      dtype=tf.float32,
                                      name='reference_shape')

        image_shape = _images[0].shape
        lms_shape = _shapes[0].points.shape

        def get_random_sample(rotation_stddev=10):
            idx = np.random.randint(low=0, high=len(_images))
            im = menpo.image.Image(_images[idx].transpose(2, 0, 1), copy=False)
            lms = _shapes[idx]
            im.landmarks['PTS'] = lms

            if np.random.rand() < .5:
               im = utils.mirror_image(im)

            if np.random.rand() < .5:
              theta = np.random.normal(scale=rotation_stddev)
              rot = menpo.transform.rotate_ccw_about_centre(lms, theta)
              im = im.warp_to_shape(im.shape, rot)

            pixels = im.pixels.transpose(1, 2, 0).astype('float32')
            shape = im.landmarks['PTS'].lms.points.astype('float32')
            return pixels, shape

        image, shape = tf.py_func(get_random_sample, [],
                                  [tf.float32, tf.float32])

        initial_shape = data_provider.random_shape(shape, reference_shape,
                                                   pca_model)
        image.set_shape(image_shape)
        shape.set_shape(lms_shape)
        initial_shape.set_shape(lms_shape)

        image = data_provider.distort_color(image)

        images, lms, inits = tf.train.batch([image, shape, initial_shape],
                                            FLAGS.batch_size,
                                            dynamic_pad=False,
                                            capacity=5000,
                                            enqueue_many=False,
                                            num_threads=num_preprocess_threads,
                                            name='batch')
        print('Defining model...')
        with tf.device(FLAGS.train_device):
            # Retain the summaries from the final tower.
            summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)
            predictions, dxs, _ = mdm_model.model(images, inits, patch_shape=(FLAGS.patch_size, FLAGS.patch_size))

            total_loss = 0

            for i, dx in enumerate(dxs):
                norm_error = mdm_model.normalized_rmse(dx + inits, lms)
                tf.histogram_summary('errors', norm_error)
                loss = tf.reduce_mean(norm_error)
                total_loss += loss
                summaries.append(tf.scalar_summary('losses/step_{}'.format(i),
                                                   loss))

            # Calculate the gradients for the batch of data
            grads = opt.compute_gradients(total_loss)

        summaries.append(tf.scalar_summary('losses/total', total_loss))
        pred_images, = tf.py_func(utils.batch_draw_landmarks,
                                  [images, predictions], [tf.float32])
        gt_images, = tf.py_func(utils.batch_draw_landmarks, [images, lms],
                                [tf.float32])

        summary = tf.image_summary('images',
                                   tf.concat(2, [gt_images, pred_images]),
                                   max_images=5)
        summaries.append(tf.histogram_summary('dx', predictions - inits))

        summaries.append(summary)

        batchnorm_updates = tf.get_collection(slim.ops.UPDATE_OPS_COLLECTION,
                                              scope)

        # Add a summary to track the learning rate.
        summaries.append(tf.scalar_summary('learning_rate', lr))

        # Add histograms for gradients.
        for grad, var in grads:
            if grad is not None:
                summaries.append(tf.histogram_summary(var.op.name +
                                                      '/gradients', grad))

        # Apply the gradients to adjust the shared variables.
        apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)

        # Add histograms for trainable variables.
        for var in tf.trainable_variables():
            summaries.append(tf.histogram_summary(var.op.name, var))

        # Track the moving averages of all trainable variables.
        # Note that we maintain a "double-average" of the BatchNormalization
        # global statistics. This is more complicated then need be but we employ
        # this for backward-compatibility with our previous models.
        variable_averages = tf.train.ExponentialMovingAverage(
            MOVING_AVERAGE_DECAY, global_step)

        # Another possibility is to use tf.slim.get_variables().
        variables_to_average = (
            tf.trainable_variables() + tf.moving_average_variables())
        variables_averages_op = variable_averages.apply(variables_to_average)

        # Group all updates to into a single train op.
        # NOTE: Currently we are not using batchnorm in MDM.
        batchnorm_updates_op = tf.group(*batchnorm_updates)
        train_op = tf.group(apply_gradient_op, variables_averages_op,
                            batchnorm_updates_op)

        # Create a saver.
        saver = tf.train.Saver(tf.all_variables())

        # Build the summary operation from the last tower summaries.
        summary_op = tf.merge_summary(summaries)
        # Start running operations on the Graph. allow_soft_placement must be
        # set to True to build towers on GPU, as some of the ops do not have GPU
        # implementations.
        sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
        # Build an initialization operation to run below.
        init = tf.initialize_all_variables()
        print('Initializing variables...')
        sess.run(init)
        print('Initialized variables.')

        step = 1
        pretrainedPath = FLAGS.pretrained_model_checkpoint_path
        if FLAGS.pretrained_model_step != '':
            pretrainedPath = pretrainedPath + '-' + str(FLAGS.pretrained_model_step)
            step = FLAGS.pretrained_model_step + 1

        if FLAGS.pretrained_model_checkpoint_path:
            assert tf.gfile.Exists(pretrainedPath)
            variables_to_restore = tf.get_collection(
                slim.variables.VARIABLES_TO_RESTORE)
            restorer = tf.train.Saver(variables_to_restore)
            restorer.restore(sess, pretrainedPath)
            print('%s: Pre-trained model restored from %s' %
                  (datetime.now(), pretrainedPath))

        # Start the queue runners.
        tf.train.start_queue_runners(sess=sess)

        summary_writer = tf.train.SummaryWriter(FLAGS.train_dir)

        print('Starting training...')
        while step < FLAGS.max_steps:
            start_time = time.time()
            _, loss_value = sess.run([train_op, total_loss])
            duration = time.time() - start_time

            assert not np.isnan(loss_value), 'Model diverged with loss = NaN'

            if step % 1 == 0:
                examples_per_sec = FLAGS.batch_size / float(duration)
                format_str = (
                    '%s: step %d, loss = %.2f (%.1f examples/sec; %.3f '
                    'sec/batch)')
                print(format_str % (datetime.now(), step, loss_value,
                                    examples_per_sec, duration))

            if step % 300 == 0:
                summary_str = sess.run(summary_op)
                summary_writer.add_summary(summary_str, step)

            # Save the model checkpoint periodically.
            if step % 300 == 0 or (step + 1) == FLAGS.max_steps:
                checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
                saver.save(sess, checkpoint_path, global_step=step)

            step = step + 1

if __name__ == '__main__':
    train()


# from __future__ import absolute_import

import torch
import torch.nn as nn
from torch.autograd import Variable

import math
import copy
import numpy as np
from . import resnet_res4s1
from . import inflated_resnet
import torchvision

import torch.nn.functional as F
from geotnf.transformation import GeometricTnfAffine
from geotnf.loss import TransformedGridLoss, WeakInlierCountPool
from utils.torch_util import expand_dim

import random
import utils.imutils2

import time
import sys

class CycleTime(nn.Module):

    def __init__(self, class_num=8, dim_in=2048, trans_param_num=3, detach_network=False, pretrained=True, temporal_out=4, T=None, hist=1):
        super(CycleTime, self).__init__()

        dim = 512
        #print("Pretrained imagenet:", pretrained)

        resnet = resnet_res4s1.resnet18(pretrained=pretrained)
        #print("RESNET 18:", resnet)

        #resnet_50 = resnet_res4s1.resnet50(pretrained=pretrained)
        #print("RESNET 50:", resnet_50)

        self.encoderVideo = inflated_resnet.InflatedResNet(copy.deepcopy(resnet))
        #print("INFLATED RESNET 18:", self.encoderVideo)

        self.detach_network = detach_network
        self.hist = hist

        self.div_num = 512
        self.T = self.div_num**-.5 if T is None else T
        print('self.T:', self.T)

        self.afterconv1 = nn.Conv3d(256, 18, kernel_size=1, bias=False)

        self.spatial_out1 = 30
        self.spatial_out2 = 10
        self.temporal_out = temporal_out

        self.afterconv3_trans = nn.Conv2d(self.spatial_out1 * self.spatial_out1, 128, kernel_size=4, padding=0, bias=False)
        self.afterconv4_trans = nn.Conv2d(128, 64, kernel_size=4, padding=0, bias=False)

        corrdim = 64 * 4 * 4
        corrdim_trans = 64 * 4 * 4

        self.linear2 = nn.Linear(corrdim_trans, trans_param_num)

        self.leakyrelu = nn.LeakyReLU(0.1, inplace=True)
        self.relu = nn.ReLU(inplace=True)
        self.avgpool = nn.AvgPool2d(7, stride=1)

        self.avgpool3d = nn.AvgPool3d((4, 2, 2), stride=(1, 2, 2))
        self.maxpool2d = nn.MaxPool2d(2, stride=2)

        # initialization

        nn.init.kaiming_normal(self.afterconv1.weight, mode='fan_out')
        nn.init.kaiming_normal(self.afterconv3_trans.weight, mode='fan_out')
        nn.init.kaiming_normal(self.afterconv4_trans.weight, mode='fan_out')

        # assuming no fc pre-training
        for m in self.modules():
            if isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

        # transformation
        self.geometricTnf = GeometricTnfAffine(geometric_model='affine',
                                         tps_grid_size=3,
                                         tps_reg_factor=0.2,
                                         out_h=self.spatial_out2, out_w=self.spatial_out2,
                                         offset_factor=227/210)

        xs = np.linspace(-1,1,80)
        xs = np.meshgrid(xs, xs)
        xs = np.stack(xs, 2)
        self.xs = xs

        self.criterion_inlier = WeakInlierCountPool(geometric_model='affine', tps_grid_size=3, tps_reg_factor=0.2, h_matches=30, w_matches=30, use_conv_filter=False, dilation_filter=0, normalize_inlier_count=True)
        self.criterion_synth  = TransformedGridLoss(use_cuda=True, geometric_model='affine')

    def compute_corr_softmax(self, patch_feat1, r50_feat2, detach_corrfeat=False):
        T = r50_feat2.shape[2]

        if detach_corrfeat is True:
            r50_feat2 = r50_feat2.detach()

        r50_feat2 = r50_feat2.transpose(3, 4) # for the inlier counter
        r50_feat2 = r50_feat2.contiguous()
        r50_feat2_vec = r50_feat2.view(r50_feat2.size(0), r50_feat2.size(1), -1)
        r50_feat2_vec = r50_feat2_vec.transpose(1, 2)

        patch_feat1_vec = patch_feat1.view(patch_feat1.size(0), patch_feat1.size(1), -1)
        corrfeat = torch.matmul(r50_feat2_vec, patch_feat1_vec)

        corrfeat = torch.div(corrfeat, self.T)

        corrfeat  = corrfeat.view(corrfeat.size(0), T, self.spatial_out1 * self.spatial_out1, self.spatial_out2, self.spatial_out2)
        corrfeat  = F.softmax(corrfeat, dim=2)
        corrfeat  = corrfeat.view(corrfeat.size(0), T * self.spatial_out1 * self.spatial_out1, self.spatial_out2, self.spatial_out2)

        return corrfeat
    #
    def compute_corr_softmax2(self, patch_feat1, r50_feat2):
        T = r50_feat2.shape[2]

        # bs, channels, time, h2, w2
        r50_feat2 = r50_feat2.contiguous()
        r50_feat2_vec = r50_feat2.view(r50_feat2.size(0), r50_feat2.size(1), -1)
        r50_feat2_vec = r50_feat2_vec.transpose(1, 2)

        # bs, channels, h1, w1
        patch_feat1 = patch_feat1.transpose(2, 3)
        patch_feat1 = patch_feat1.contiguous()
        patch_feat1_vec = patch_feat1.view(patch_feat1.size(0), patch_feat1.size(1), -1)

        corrfeat = torch.matmul(r50_feat2_vec, patch_feat1_vec)
        corrfeat  = torch.div(corrfeat, self.T)
        corrfeat  = corrfeat.contiguous()

        corrfeat  = corrfeat.view(corrfeat.size(0), T, self.spatial_out2 * self.spatial_out2, self.spatial_out1 * self.spatial_out1)
        corrfeat  = F.softmax(corrfeat, dim=3)
        corrfeat  = corrfeat.transpose(2, 3)
        corrfeat  = corrfeat.contiguous()
        corrfeat  = corrfeat.view(corrfeat.size(0) * T, self.spatial_out1 * self.spatial_out1, self.spatial_out2, self.spatial_out2)

        return corrfeat

    def transform_trans_out(self, trans_out1):

        trans_out1 = trans_out1.view(-1, 3)

        trans_out1_theta = trans_out1[:, 2]
        trans_out1_2 = trans_out1[:, 0].unsqueeze(1)
        trans_out1_5 = trans_out1[:, 1].unsqueeze(1)
        trans_out1_0 = 1.0 / 3.0 * torch.cos(trans_out1_theta).unsqueeze(1)
        trans_out1_1 = - 1.0 / 3.0 * torch.sin(trans_out1_theta).unsqueeze(1)
        trans_out1_3 = 1.0 / 3.0 * torch.sin(trans_out1_theta).unsqueeze(1)
        trans_out1_4 = 1.0 / 3.0 * torch.cos(trans_out1_theta).unsqueeze(1)
        trans_out1 = torch.cat((trans_out1_0, trans_out1_1, trans_out1_2, trans_out1_3, trans_out1_4, trans_out1_5), dim=1)

        trans_out1 = trans_out1.view(-1, 2, 3)

        return trans_out1

    def forward_base(self, x, contiguous=False, can_detach=True):
        # import pdb; pdb.set_trace()

        # patch feature
        x = x.transpose(1, 2)
        x_pre = self.encoderVideo(x)

        if self.detach_network and can_detach:
            x_pre = x_pre.detach()

        x = self.afterconv1(x_pre)
        x = self.relu(x)

        if contiguous:
            x = x.contiguous()
            x = x.view(x.size(0), x.size(1), x.size(3), x.size(4))

        x_norm = F.normalize(x, p=2, dim=1)

        return x, x_pre, x_norm

    def compute_transform_img_to_patch(self, query, base, temporal_out=1, detach_corrfeat=False):

        # query is a patch, base is a volume with time depth 1
        corrfeat = self.compute_corr_softmax(query, base, detach_corrfeat=detach_corrfeat)

        # 2. predict transform with affinity as input
        corrfeat_mat = corrfeat.view(corrfeat.size(0) * temporal_out, self.spatial_out1 * self.spatial_out1, self.spatial_out2, self.spatial_out2)
        corrfeat_trans  = self.afterconv3_trans(corrfeat_mat)
        corrfeat_trans  = self.leakyrelu(corrfeat_trans)
        corrfeat_trans  = self.afterconv4_trans(corrfeat_trans)
        corrfeat_trans  = self.leakyrelu(corrfeat_trans)
        corrfeat_trans  = corrfeat_trans.view(corrfeat_trans.shape[0], -1)

        trans_theta = self.linear2(corrfeat_trans).contiguous().view(corrfeat_trans.shape[0], 3)
        trans_theta = self.transform_trans_out(trans_theta)

        return corrfeat, corrfeat_mat, corrfeat_trans, trans_theta

    def forward(self, ximg1, patch2, img2, theta):
        B, T = ximg1.shape[:2]
        videoclip1  = ximg1

        # base features
        r50_feat1, r50_feat1_pre, r50_feat1_norm = self.forward_base(videoclip1)

        # target patch feature
        patch2_feat2, patch2_feat2_pre, patch_feat2_norm = self.forward_base(patch2, contiguous=True)

        # target image feature
        img_feat2, img_feat2_pre, img_feat2_norm = self.forward_base(img2, contiguous=True, can_detach=False)

        # base features to crop with transformations
        r50_feat1_transform = r50_feat1.transpose(1, 2)
        channels = r50_feat1_transform.size(2)
        r50_feat1_transform = r50_feat1_transform.contiguous()

        # add original code
        corrfeat1, corrfeat_trans_matrix2, corrfeat_trans1, trans_out2 = self.compute_transform_img_to_patch(patch_feat2_norm, r50_feat1_norm, temporal_out=self.temporal_out)

        bs2 = corrfeat_trans1.size(0)

        r50_feat1_transform_ori = r50_feat1_transform.view(bs2, channels, self.spatial_out1, self.spatial_out1)
        r50_feat1_transform_ori = self.geometricTnf(r50_feat1_transform_ori, trans_out2)

        # r50_feat1_transform_ori = r50_feat1_transform_ori.transpose(1, 2)

        def skip_prediction(img_feat2_norm, r50_feat1_transform_ori):
            r50_feat1_transform_ori = r50_feat1_transform_ori.contiguous()
            r50_feat1_transform_ori = r50_feat1_transform_ori.view(B, self.temporal_out, r50_feat1_transform_ori.size(1),  self.spatial_out2, self.spatial_out2)
            r50_feat1_transform_ori = r50_feat1_transform_ori.transpose(1, 2)

            r50_feat1_transform_norm = F.normalize(r50_feat1_transform_ori, p=2, dim=1)
            corrfeat_trans_matrix_reverse = self.compute_corr_softmax2(img_feat2_norm, r50_feat1_transform_norm)

            corrfeat_trans_reverse  = self.afterconv3_trans(corrfeat_trans_matrix_reverse)
            corrfeat_trans_reverse  = self.leakyrelu(corrfeat_trans_reverse)
            corrfeat_trans_reverse  = self.afterconv4_trans(corrfeat_trans_reverse)
            corrfeat_trans_reverse  = self.leakyrelu(corrfeat_trans_reverse)
            corrfeat_trans_reverse  = corrfeat_trans_reverse.contiguous()
            corrfeat_trans_reverse  = corrfeat_trans_reverse.view(bs2, -1)

            trans_out3  = self.linear2(corrfeat_trans_reverse)
            trans_out3  = trans_out3.contiguous()
            trans_out3  = self.transform_trans_out(trans_out3)

            return trans_out3, corrfeat_trans_matrix_reverse

        trans_out3, corrfeat_trans_matrix_reverse = skip_prediction(img_feat2_norm, r50_feat1_transform_ori)

        def recurrent_align(init_query, idx):
            # global ximg1
            # global patch2

            corr_feat_mats = []
            trans_thetas = []
            trans_feats = []

            # should be normalized patch query
            if self.hist > 1:
                cur_query = torch.stack([init_query]*self.hist)
            else:
                cur_query = init_query

            crops = []
            for t in idx:

                # 1. get affinity of current patch on current frame
                cur_base_norm = r50_feat1_norm[:, :, t:t+1]
                cur_base_feat = r50_feat1_transform[:, t]

                # 2. predict transform with affinity as input
                corrfeat, corrfeat_mat, corrfeat_trans, trans_theta = self.compute_transform_img_to_patch(
                    cur_query if self.hist == 1 else cur_query.mean(0),
                    cur_base_norm)

                # 3. get cropped features with transform
                cur_base_crop = self.geometricTnf(cur_base_feat, trans_theta)

                # bs, channels, time, h2, w2
                cur_base_crop_norm = F.normalize(cur_base_crop, p=2, dim=1)

                # cur_query = cur_base_crop_norm

                if self.hist > 1:
                    cur_query[:-1] = cur_query[1:]
                    cur_query[-1] = cur_base_crop_norm
                else:
                    cur_query = cur_base_crop_norm
                # cur_query = torch.stack([cur_])

                trans_thetas.append(trans_theta)
                trans_feats.append(cur_base_crop)
                # corr_feat_mats.append(corrfeat_mat)

            return trans_thetas, trans_feats #, trans_feats, corr_feat_mats

        def cycle(TT=None):
            if TT is None:
                TT = T

            # propagate backward
            # back_trans_thetas, back_trans_feats, back_corr_feat_mats = \
            back_trans_thetas, back_trans_feats = \
                recurrent_align(patch_feat2_norm, list(range(T))[::-1][:TT])
            # propagate forward
            # forw_trans_thetas, forw_trans_feats, forw_corr_feat_mats = \
            forw_trans_thetas, forw_trans_feats = \
                recurrent_align(F.normalize(back_trans_feats[-1], p=2, dim=1), list(range(T))[T-TT+1:])

            # cycle back from last base frame to target
            last_ = forw_trans_feats[-1] if len(forw_trans_feats) > 0 else back_trans_feats[0]
            last_corrfeat, last_corrfeat_mat, last_corrfeat_trans, last_trans_theta = self.compute_transform_img_to_patch(
                F.normalize(last_, p=2, dim=1), img_feat2_norm.unsqueeze(2))
            last_trans_feat = self.geometricTnf(img_feat2, last_trans_theta)
            last_trans_feat_norm = F.normalize(last_trans_feat, p=2, dim=1)

            forw_trans_thetas.append(last_trans_theta)

            return back_trans_thetas, forw_trans_thetas, back_trans_feats

        # back_trans_thetas, back_trans_feats, back_corr_feat_mats, forw_trans_thetas, forw_trans_feats, forw_corr_feat_mats = \
        #     [], [], [], [], [], []

        outputs = [[], [], []]

        for c in range(1, T+1):
            # _back_trans_thetas, _back_trans_feats, _back_corr_feat_mats, _forw_trans_thetas, _forw_trans_feats, _forw_corr_feat_mats = cycle(c)
            _outputs = cycle(c)
            for i, o in enumerate(_outputs):
                outputs[i] += o

            if c == T:
                back_trans_feats = _outputs[-1]

        back_trans_feats = torch.stack(back_trans_feats).transpose(0,1).contiguous()
        back_trans_feats = back_trans_feats.view(-1, *back_trans_feats.shape[2:])
        skip_trans, skip_corrfeat_mat = skip_prediction(img_feat2_norm, back_trans_feats)

        return outputs[:2], patch2_feat2, theta, trans_out2, trans_out3, skip_trans, skip_corrfeat_mat, corrfeat_trans_matrix2

    def loss(self, outputs, patch_feat, theta, trans_out2, trans_out3, skip_trans, skip_corrfeat_mat, corrfeat_trans_matrix2):
        # patch_feat is patch of target frame, theta is crop transform for patch

        forw_trans_thetas, back_trans_thetas = outputs

        loss_targ_theta = []
        loss_targ_theta_skip = []
        loss_back_inliers = []

        nn = list(range(len(forw_trans_thetas)))
        nn = [ii for ii in [sum(nn[:i]) - 1 for i in nn][2:] if ii < len(forw_trans_thetas)]

        for i in nn:
            loss_targ_theta.append(self.criterion_synth(forw_trans_thetas[i], theta))

        theta2 = theta.unsqueeze(1)
        theta2 = theta2.repeat(1, self.temporal_out, 1, 1)
        theta2 = theta2.view(-1, 2, 3)

        loss_targ_theta_skip.append(self.criterion_synth(trans_out3, theta2))

        loss_inlier = self.criterion_inlier(matches=corrfeat_trans_matrix2, theta=trans_out2)
        loss_inlier = torch.mean(-loss_inlier)
        loss_back_inliers.append(loss_inlier)

        return (loss_targ_theta, loss_targ_theta_skip, loss_back_inliers)


#!/usr/bin/env python
# -*- coding=utf-8 -*-
###########################################################################
# Copyright (C) 2013-2021 by Caspar. All rights reserved.
# File Name: .py
# Author: Shankai Yan
# E-mail: dr.skyan@gmail.com
# Created Time: 2021-03-29 17:19:30
###########################################################################
#

import numpy as np
from scipy.sparse import csc_matrix

import torch
from torch import nn
import torch.nn.functional as F

from . import transformer as T

class SiameseRankHead(nn.Module):
    def __init__(self, linear, siamese_linear, binlbr, tokenizer, encode_func, trnsfm, special_tknids_args, pad_val, base_model, thrshld_lnr=0.5, thrshld_sim=0.5, topk=None, num_sampling_note=1, lbnotes='lbnotes.csv'):
        super(SiameseRankHead, self).__init__()
        self.linear = linear
        self.siamese_linear = siamese_linear
        self.base_model = dict(zip(['model', 'tokenizer', 'encode_func', 'special_tknids_args', 'pad_val', 'transforms', 'transforms_args', 'transforms_kwargs'], [base_model, tokenizer, encode_func, special_tknids_args, pad_val]+trnsfm))
        self.binlbr = binlbr
        self.thrshld_lnr = thrshld_lnr
        self.thrshld_sim = thrshld_sim
        self.topk = topk
        self.num_sampling_note = num_sampling_note
        self.lbnotes = lbnotes if type(lbnotes) is pd.DataFrame else pd.read_csv(lbnotes, sep='\t', index_col='id', encoding='utf-8')

    def forward(self, hidden_states):
        # generate candidates with linear and rank the candidates with siamese_linear
        use_gpu = next(self.parameters()).is_cuda
        batch_size = 125
        logits = self.linear(hidden_states)
        prob = torch.sigmoid(logits)
        if self.topk:
            topk, topk_indices = torch.topk(prob.data, self.topk)
            prob_topk = torch.zeros(prob.size(), dtype=topk.dtype).scatter_(1, topk_indices.cpu(), topk.cpu())
            pred = (prob_topk > self.thrshld_lnr).int()
        else:
            pred = (prob.data > self.thrshld_lnr).int()
        pred_csc = csc_matrix(pred.cpu().numpy())
        pred_scores = [[] for x in range(len(pred_csc.data))]
        labels = [self.binlbr[i] for i in range(len(self.binlbr))]
        num_notes = np.array([min(self.num_sampling_note, self.lbnotes.loc[lb].shape[0]) if lb in self.lbnotes.index else 0 for lb in labels])
        num_samps = pred_csc.indptr[1:] - pred_csc.indptr[:-1]
        total_num_notes, total_num_pairs = sum(num_notes), sum(num_samps * num_notes)
        def gen_samples():
            for i, lb in enumerate(labels):
                for note in (self.lbnotes.loc[lb]['text'].sample(n=min(self.num_sampling_note, self.lbnotes.loc[lb].shape[0])) if type(self.lbnotes.loc[lb]['text']) is pd.Series else [self.lbnotes.loc[lb]['text']]) if lb in self.lbnotes.index else []:
                    tkns = self._transform_chain((self.base_model['encode_func'](note, self.base_model['tokenizer']), None), self.base_model['transforms'], self.base_model['transforms_args'], self.base_model['transforms_kwargs'])[0]
                    yield tkns, i, pred_csc.indices[pred_csc.indptr[i]:pred_csc.indptr[i+1]]
        def gen_pairs():
            batch_tkns, lb_indices, orig_indices = [[] for x in range(3)]
            for i, (x, lbidx, orig_idx) in enumerate(gen_samples()):
                batch_tkns.append(x)
                lb_indices.append(lbidx)
                orig_indices.append(orig_idx)
                if i > 0 and i % batch_size == 0 or i == total_num_notes - 1:
                    clf_tknids = self.base_model['special_tknids_args']['clf_tknids']
                    tkns_tnsr = torch.tensor(batch_tkns, dtype=torch.long).to('cuda') if use_gpu else torch.tensor(batch_tkns, dtype=torch.long)
                    mask_tnsr = (~tkns_tnsr.eq(self.base_model['pad_val'] * torch.ones_like(tkns_tnsr))).long()
                    pool_idx = tkns_tnsr.eq(clf_tknids[0] * torch.ones_like(tkns_tnsr)).int().argmax(-1)
                    note_clf_h = self.base_model['model'](tkns_tnsr, mask_tnsr if type(self.base_model['model']) is T.BERTClfHead else pool_idx, embedding_mode=True)
                    # note_clf_h = hidden_states[:1]
                    for j, clf_h in enumerate(note_clf_h):
                        for orig_id in orig_indices[j]:
                            yield clf_h, hidden_states[orig_id], lb_indices[j], orig_id
                    batch_tkns, lb_indices, orig_indices = [[] for x in range(3)]
        batch_pairs, indices = [[] for x in range(2)]
        for i, (h1, h2, lbidx, orig_id) in enumerate(tqdm(gen_pairs(), desc='[%i pair(s)] Siamese similarity ranking' % total_num_pairs)):
            batch_pairs.append([h1, h2])
            indices.append([lbidx, orig_id])
            if i > 0 and i % batch_size == 0 or i == total_num_pairs - 1:
                clf_h = [torch.stack(x) for x in zip(*batch_pairs)]
                continue
                if (self.base_model['model'].task_params.setdefault('sentsim_func', None) is None or self.base_model['model'].task_params['sentsim_func'] == 'concat'):
                    clf_h = torch.cat(clf_h, dim=-1)
                    sim_logits = self.siamese_linear(clf_h) if self.siamese_linear else clf_h
                else:
                    sim_logits = F.pairwise_distance(self.siamese_linear(clf_h[0]), self.siamese_linear(clf_h[1]), 2, eps=1e-12) if self.base_model['model'].task_params['sentsim_func'] == 'dist' else F.cosine_similarity(self.siamese_linear(clf_h[0]), self.siamese_linear(clf_h[1]), dim=1, eps=1e-12)
                sim_probs = torch.sigmoid(sim_logits).data
                for idx, sim_prob in zip(indices, sim_probs):
                    orig_idx = pred_csc.indptr[idx[0]], pred_csc.indptr[idx[0]+1]
                    pred_scores[orig_idx[0]:orig_idx[1]][np.where(pred_csc.indices[orig_idx[0]:orig_idx[1]]==idx[1])[0][0]].append(sim_prob.cpu().numpy())
                batch_pairs, indices = [[] for x in range(2)]
        pred_scores = [np.mean(x) for x in pred_scores]
        sim_probs_tnsr = torch.tensor(csc_matrix((pred_scores, pred_csc.indices, pred_csc.indptr), shape=pred_csc.shape).todense(), dtype=torch.float)
        if use_gpu: sim_probs_tnsr = sim_probs_tnsr.to('cuda')
        logits = sim_probs_tnsr * prob
        return logits

    def to(self, *args, **kwargs):
        super(SiameseRankHead, self).to(*args, **kwargs)
        self.linear = self.linear.to(*args, **kwargs)
        return self

    def _transform_chain(self, sample, transforms, transforms_args={}, transforms_kwargs={}):
        if transforms:
            transforms = transforms if type(transforms) is list else [transforms]
            transforms_kwargs = transforms_kwargs if type(transforms_kwargs) is list else [transforms_kwargs]
            for transform, transform_kwargs in zip(transforms, transforms_kwargs):
                transform_kwargs.update(transforms_args)
                sample = transform(sample, **transform_kwargs) if callable(transform) else getattr(self, transform)(sample, **transform_kwargs)
        return sample

class SiameseRankTransformer(object):
    def __init__(self, clf):
        self.clf = clf

    def merge_siamese(self, tokenizer, encode_func, trnsfm, special_tknids_args, pad_val=0, thrshld_lnr=0.5, thrshld_sim=0.5, topk=None, lbnotes='lbnotes.csv'):
        use_gpu = next(self.clf.parameters()).is_cuda
        self.clf.task_type = self.clf.clf_task_type
        self.clf.linear = SiameseRankHead(self.clf.clf_linear, self.clf.linear, self.clf.binlbr, tokenizer=tokenizer, encode_func=encode_func, trnsfm=trnsfm, special_tknids_args=special_tknids_args, pad_val=pad_val, base_model=self.clf, thrshld_lnr=thrshld_lnr, thrshld_sim=thrshld_sim, topk=topk, lbnotes=lbnotes)
        self.clf.linear = self.clf.linear.to('cuda') if use_gpu else self.clf.linear
        self.clf.mode = 'clf'


# @manual=//deeplearning/trt/python:py_tensorrt
import numpy as np
import tensorrt as trt
import torch

from ..converter_registry import tensorrt_converter

from .converter_utils import (
    extend_mod_attr_to_tuple,
    get_dyn_range,
    mark_as_int8_layer,
    to_numpy,
)

def common_conv(network, mod, dimension, input_val, layer_name, is_quantized):
    if mod.padding_mode != "zeros":
        raise RuntimeError(f"Only support padding mode: zeros, got {mod.padding_mode}.")

    kernel_size = extend_mod_attr_to_tuple(mod, "kernel_size", dimension)
    stride = extend_mod_attr_to_tuple(mod, "stride", dimension)
    padding = extend_mod_attr_to_tuple(mod, "padding", dimension)
    dilation = extend_mod_attr_to_tuple(mod, "dilation", dimension)

    kernel = to_numpy(mod.weight() if is_quantized else mod.weight)
    bias = to_numpy(mod.bias() if is_quantized else mod.bias)

    if dimension == 1:
        # Append unsqueeze before conv2d to calculate conv1d
        unsqueeze_layer = network.add_shuffle(input=input_val)
        unsqueeze_layer.reshape_dims = (*input_val.shape, 1)
        unsqueeze_layer.name = f"{layer_name}_unsqueeze"
        input_val = unsqueeze_layer.get_output(0)

        kernel = np.expand_dims(kernel, -1)
        kernel_size = kernel.shape[2:]
        if bias is not None:
            bias = bias[None]
        stride = (stride[0], 1)
        padding = (padding[0], 0)
        dilation = (dilation[0], 1)
    layer = network.add_convolution_nd(
        input=input_val,
        num_output_maps=mod.out_channels,
        kernel_shape=kernel_size,
        kernel=kernel,
        bias=bias,
    )
    layer.name = layer_name
    layer.stride_nd = stride
    layer.padding_nd = padding
    layer.dilation_nd = dilation
    layer.num_groups = mod.groups

    if is_quantized:
        # Assume the dtype of activation is torch.quint8
        mark_as_int8_layer(
            layer, get_dyn_range(mod.scale, mod.zero_point, torch.quint8)
        )

    result = layer.get_output(0)
    if dimension == 1:
        # Append squeeze after conv2d to calculate conv1d
        squeeze_layer = network.add_shuffle(input=result)
        squeeze_layer.reshape_dims = tuple(result.shape[:-1])
        squeeze_layer.name = f"{layer_name}_squeeze"
        result = squeeze_layer.get_output(0)

    return result

def common_conv_relu(network, mod, dimension, input_val, layer_name, is_quantized):
    conv_output = common_conv(
        network,
        mod,
        dimension=2,
        input_val=input_val,
        layer_name=f"{layer_name}_conv",
        is_quantized=is_quantized,
    )

    layer = network.add_activation(input=conv_output, type=trt.ActivationType.RELU)
    layer.name = f"{layer_name}_relu"

    if is_quantized:
        mark_as_int8_layer(layer, conv_output.dynamic_range)

    return layer.get_output(0)

@tensorrt_converter(torch.nn.modules.conv.Conv1d)
def conv1d(network, submod, args, kwargs, layer_name):
    # args/kwargs should have already been normalized to kwargs
    assert len(args) == 0
    input_val = kwargs["input"]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Conv1d received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    if layer_name is None:
        raise RuntimeError("layer name is none")
    return common_conv(
        network,
        submod,
        dimension=1,
        input_val=input_val,
        layer_name=layer_name,
        is_quantized=False,
    )

@tensorrt_converter(torch.nn.modules.conv.Conv2d)
def conv2d(network, submod, args, kwargs, layer_name):
    # args/kwargs should have already been normalized to kwargs
    assert len(args) == 0
    input_val = kwargs["input"]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Conv2d received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    return common_conv(
        network,
        submod,
        dimension=2,
        input_val=input_val,
        layer_name=layer_name,
        is_quantized=False,
    )

@tensorrt_converter(torch.nn.modules.conv.Conv3d)
def conv3d(network, submod, args, kwargs, layer_name):
    # args/kwargs should have already been normalized to kwargs
    assert len(args) == 0
    input_val = kwargs["input"]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Conv3d received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    return common_conv(
        network,
        submod,
        dimension=3,
        input_val=input_val,
        layer_name=layer_name,
        is_quantized=False,
    )

@tensorrt_converter(torch.nn.quantized.modules.conv.Conv2d)
def quantized_conv2d(network, submod, args, kwargs, layer_name):
    input_val = args[0]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Quantized Conv2d received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    return common_conv(
        network,
        submod,
        dimension=2,
        input_val=input_val,
        layer_name=layer_name,
        is_quantized=True,
    )

@tensorrt_converter(torch.nn.intrinsic.quantized.modules.ConvReLU2d)
def quantized_conv_relu2d(network, submod, args, kwargs, layer_name):
    input_val = args[0]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Quantized ConvReLU2d received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    return common_conv_relu(
        network,
        submod,
        dimension=2,
        input_val=input_val,
        layer_name=f"{layer_name}_conv",
        is_quantized=True,
    )