# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
This module provides some useful functions for dealing with magnetic Structures
(e.g. Structures with associated magmom tags).
"""

import warnings
import numpy as np
import os
import logging

from enum import Enum, unique
from collections import namedtuple

from scipy.stats import gaussian_kde
from scipy.signal import argrelextrema

from pymatgen.core.structure import Species, Structure, Element, DummySpecies
from pymatgen.electronic_structure.core import Magmom
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.transformations.standard_transformations import (
    AutoOxiStateDecorationTransformation,
)
from pymatgen.transformations.advanced_transformations import (
    MagOrderingTransformation,
    MagOrderParameterConstraint,
)
from pymatgen.symmetry.groups import SpaceGroup
from monty.serialization import loadfn

from typing import Union, List, Dict, Tuple, Optional, Any
from pymatgen.util.typing import Vector3Like

"""
This module provides some useful functions for dealing with magnetic Structures
(e.g. Structures with associated magmom tags).
"""

__author__ = "Matthew Horton"
__copyright__ = "Copyright 2017, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Matthew Horton"
__email__ = "mkhorton@lbl.gov"
__status__ = "Development"
__date__ = "Feb 2017"

MODULE_DIR = os.path.dirname(os.path.abspath(__file__))

try:
    DEFAULT_MAGMOMS = loadfn(os.path.join(MODULE_DIR, "default_magmoms.yaml"))
except Exception:
    warnings.warn(
        "Could not load default_magmoms.yaml, falling back to VASPIncarBase.yaml"
    )
    DEFAULT_MAGMOMS = loadfn(
        os.path.join(MODULE_DIR, "../../io/vasp/VASPIncarBase.yaml")
    )
    DEFAULT_MAGMOMS = DEFAULT_MAGMOMS["MAGMOM"]

@unique
class Ordering(Enum):
    """
    Enumeration defining possible magnetic orderings.
    """

    FM = "FM"  # Ferromagnetic
    AFM = "AFM"  # Antiferromagnetic
    FiM = "FiM"  # Ferrimagnetic
    NM = "NM"  # Non-magnetic
    Unknown = "Unknown"

@unique
class OverwriteMagmomMode(Enum):
    """
    Enumeration defining different modes for analyzer.
    """

    none = "none"
    respect_sign = "respect_sign"
    respect_zero = "respect_zeros"
    replace_all = "replace_all"
    normalize = "normalize"

class CollinearMagneticStructureAnalyzer:
    """
    A class which provides a few helpful methods to analyze
    collinear magnetic structures.
    """
    def __init__(
        self,
        structure: Structure,
        overwrite_magmom_mode: Union[OverwriteMagmomMode, str] = "none",
        round_magmoms: bool = False,
        detect_valences: bool = False,
        make_primitive: bool = True,
        default_magmoms: dict = None,
        set_net_positive: bool = True,
        threshold: float = 0.00,
        threshold_nonmag: float = 0.1,
    ):
        r"""
        If magnetic moments are not defined, moments will be
        taken either from default_magmoms.yaml (similar to the
        default magmoms in MPRelaxSet, with a few extra definitions)
        or from a specie:magmom dict provided by the default_magmoms
        kwarg.

        Input magmoms can be replaced using the 'overwrite_magmom_mode'
        kwarg. This can be:
        * "none" to do nothing,
        * "respect_sign" which will overwrite existing magmoms with
          those from default_magmoms but will keep sites with positive magmoms
          positive, negative magmoms negative and zero magmoms zero,
        * "respect_zeros", which will give a ferromagnetic structure
          (all positive magmoms from default_magmoms) but still keep sites with
          zero magmoms as zero,
        * "replace_all" which will try to guess initial magmoms for
          all sites in the structure irrespective of input structure
          (this is most suitable for an initial DFT calculation),
        * "replace_all_if_undefined" is the same as "replace_all" but only if
          no magmoms are defined in input structure, otherwise it will respect
          existing magmoms.
        * "normalize" will normalize magmoms to unity, but will respect sign
          (used for comparing orderings), magmoms < theshold will be set to zero

        Args:
            structure: input Structure object
            overwrite_magmom_mode: "respect_sign", "respect_zeros", "replace_all",
                "replace_all_if_undefined", "normalize" (default "none")
            round_magmoms: will round input magmoms to
                specified number of decimal places if integer is supplied, if set
                to a float will try and group magmoms together using a kernel density
                estimator of provided width, and extracting peaks of the estimator
                detect_valences: if True, will attempt to assign valences
                to input structure
            make_primitive: if True, will transform to primitive
                magnetic cell
            default_magmoms: (optional) dict specifying default magmoms
            set_net_positive: if True, will change sign of magnetic
                moments such that the net magnetization is positive. Argument will be
                ignored if mode "respect_sign" is used.
            threshold: number (in Bohr magnetons) below which magmoms
                will be rounded to zero
            threshold_nonmag: number (in Bohr magneton)
                below which nonmagnetic ions (with no magmom specified
                in default_magmoms) will be rounded to zero
        """

        if default_magmoms:
            self.default_magmoms = default_magmoms
        else:
            self.default_magmoms = DEFAULT_MAGMOMS

        structure = structure.copy()

        # check for disorder
        if not structure.is_ordered:
            raise NotImplementedError(
                "Not implemented for disordered structures, "
                "make ordered approximation first."
            )

        if detect_valences:
            trans = AutoOxiStateDecorationTransformation()
            try:
                structure = trans.apply_transformation(structure)
            except ValueError:
                warnings.warn(
                    "Could not assign valences "
                    "for {}".format(structure.composition.reduced_formula)
                )

        # check to see if structure has magnetic moments
        # on site properties or species spin properties,
        # prioritize site properties

        has_magmoms = bool(structure.site_properties.get("magmom", False))

        has_spin = False
        for comp in structure.species_and_occu:
            for sp, occu in comp.items():
                if getattr(sp, "spin", False):
                    has_spin = True

        # perform input sanitation ...
        # rest of class will assume magnetic moments
        # are stored on site properties:
        # this is somewhat arbitrary, arguments can
        # be made for both approaches

        if has_magmoms and has_spin:
            raise ValueError(
                "Structure contains magnetic moments on both "
                "magmom site properties and spin species "
                "properties. This is ambiguous. Remove one or "
                "the other."
            )
        elif has_magmoms:
            if None in structure.site_properties["magmom"]:
                warnings.warn(
                    "Be careful with mixing types in your magmom "
                    "site properties. Any 'None' magmoms have been "
                    "replaced with zero."
                )
            magmoms = [m if m else 0 for m in structure.site_properties["magmom"]]
        elif has_spin:
            magmoms = [getattr(sp, "spin", 0) for sp in structure.species]
            structure.remove_spin()
        else:
            # no magmoms present, add zero magmoms for now
            magmoms = [0] * len(structure)
            # and overwrite magmoms with default magmoms later unless otherwise stated
            if overwrite_magmom_mode == "replace_all_if_undefined":
                overwrite_magmom_mode = "replace_all"

        # test to see if input structure has collinear magmoms
        self.is_collinear = Magmom.are_collinear(magmoms)

        if not self.is_collinear:
            warnings.warn(
                "This class is not designed to be used with "
                "non-collinear structures. If your structure is "
                "only slightly non-collinear (e.g. canted) may still "
                "give useful results, but use with caution."
            )

        # this is for collinear structures only, make sure magmoms
        # are all floats
        magmoms = list(map(float, magmoms))

        # set properties that should be done /before/ we process input magmoms
        self.total_magmoms = sum(magmoms)
        self.magnetization = sum(magmoms) / structure.volume

        # round magmoms on magnetic ions below threshold to zero
        # and on non magnetic ions below threshold_nonmag
        magmoms = [
            m
            if abs(m) > threshold and a.species_string in self.default_magmoms
            else m
            if abs(m) > threshold_nonmag
            and a.species_string not in self.default_magmoms
            else 0
            for (m, a) in zip(magmoms, structure.sites)
        ]

        # overwrite existing magmoms with default_magmoms
        if overwrite_magmom_mode not in (
            "none",
            "respect_sign",
            "respect_zeros",
            "replace_all",
            "replace_all_if_undefined",
            "normalize",
        ):
            raise ValueError("Unsupported mode.")

        for idx, site in enumerate(structure):

            if site.species_string in self.default_magmoms:
                # look for species first, e.g. Fe2+
                default_magmom = self.default_magmoms[site.species_string]
            elif (
                isinstance(site.specie, Species)
                and str(site.specie.element) in self.default_magmoms
            ):
                # look for element, e.g. Fe
                default_magmom = self.default_magmoms[str(site.specie.element)]
            else:
                default_magmom = 0

            # overwrite_magmom_mode = "respect_sign" will change magnitude of
            # existing moments only, and keep zero magmoms as
            # zero: it will keep the magnetic ordering intact

            if overwrite_magmom_mode == "respect_sign":
                set_net_positive = False
                if magmoms[idx] > 0:
                    magmoms[idx] = default_magmom
                elif magmoms[idx] < 0:
                    magmoms[idx] = -default_magmom

            # overwrite_magmom_mode = "respect_zeros" will give a ferromagnetic
            # structure but will keep zero magmoms as zero

            elif overwrite_magmom_mode == "respect_zeros":
                if magmoms[idx] != 0:
                    magmoms[idx] = default_magmom

            # overwrite_magmom_mode = "replace_all" will ignore input magmoms
            # and give a ferromagnetic structure with magnetic
            # moments on *all* atoms it thinks could be magnetic

            elif overwrite_magmom_mode == "replace_all":
                magmoms[idx] = default_magmom

            # overwrite_magmom_mode = "normalize" set magmoms magnitude to 1

            elif overwrite_magmom_mode == "normalize":
                if magmoms[idx] != 0:
                    magmoms[idx] = int(magmoms[idx] / abs(magmoms[idx]))

        # round magmoms, used to smooth out computational data
        magmoms = (
            self._round_magmoms(magmoms, round_magmoms) if round_magmoms else magmoms
        )

        if set_net_positive:
            sign = np.sum(magmoms)
            if sign < 0:
                magmoms = -np.array(magmoms)

        structure.add_site_property("magmom", magmoms)

        if make_primitive:
            structure = structure.get_primitive_structure(use_site_props=True)

        self.structure = structure

    @staticmethod
    def _round_magmoms(
        magmoms: Vector3Like, round_magmoms_mode: Union[int, float]
    ) -> np.ndarray:
        """If round_magmoms_mode is an integer, simply round to that number
        of decimal places, else if set to a float will try and round
        intelligently by grouping magmoms.
        """

        if isinstance(round_magmoms_mode, int):

            # simple rounding to number of decimal places
            magmoms = np.around(magmoms, decimals=round_magmoms_mode)

        elif isinstance(round_magmoms_mode, float):

            try:

                # get range of possible magmoms, pad by 50% just to be safe
                range_m = max([max(magmoms), abs(min(magmoms))]) * 1.5

                # construct kde, here "round_magmoms_mode" is the width of the kde
                kernel = gaussian_kde(magmoms, bw_method=round_magmoms_mode)

                # with a linearly spaced grid 1000x finer than width
                xgrid = np.linspace(
                    -range_m, range_m, int(1000 * range_m / round_magmoms_mode)
                )

                # and evaluate the kde on this grid, extracting the maxima of the kde peaks
                kernel_m = kernel.evaluate(xgrid)
                extrema = xgrid[argrelextrema(kernel_m, comparator=np.greater)]

                # round magmoms to these extrema
                magmoms = [extrema[(np.abs(extrema - m)).argmin()] for m in magmoms]

            except Exception as e:

                # TODO: typically a singular matrix warning, investigate this
                warnings.warn(
                    "Failed to round magmoms intelligently, "
                    "falling back to simple rounding."
                )
                warnings.warn(str(e))

            # and finally round roughly to the number of significant figures in our kde width
            num_decimals = len(str(round_magmoms_mode).split(".")[1]) + 1
            magmoms = np.around(magmoms, decimals=num_decimals)

        return magmoms

    def get_structure_with_spin(self) -> Structure:
        """Returns a Structure with species decorated with spin values instead
        of using magmom site properties.
        """

        structure = self.structure.copy()
        structure.add_spin_by_site(structure.site_properties["magmom"])
        structure.remove_site_property("magmom")

        return structure

    def get_structure_with_only_magnetic_atoms(
        self, make_primitive: bool = True
    ) -> Structure:
        """Returns a Structure with only magnetic atoms present.

        Args:
          make_primitive: Whether to make structure primitive after
            removing non-magnetic atoms (Default value = True)

        Returns: Structure
        """

        sites = [site for site in self.structure if abs(site.properties["magmom"]) > 0]

        structure = Structure.from_sites(sites)

        if make_primitive:
            structure = structure.get_primitive_structure(use_site_props=True)

        return structure

    def get_nonmagnetic_structure(self, make_primitive: bool = True) -> Structure:
        """Returns a Structure without magnetic moments defined.

        Args:
          make_primitive: Whether to make structure primitive after
            removing magnetic information (Default value = True)

        Returns:
          Structure

        """

        structure = self.structure.copy()
        structure.remove_site_property("magmom")

        if make_primitive:
            structure = structure.get_primitive_structure()

        return structure

    def get_ferromagnetic_structure(self, make_primitive: bool = True) -> Structure:
        """Returns a Structure with all magnetic moments positive
        or zero.

        Args:
          make_primitive: Whether to make structure primitive after
            making all magnetic moments positive (Default value = True)

        Returns:
          Structure

        """

        structure = self.structure.copy()

        structure.add_site_property("magmom", [abs(m) for m in self.magmoms])

        if make_primitive:
            structure = structure.get_primitive_structure(use_site_props=True)

        return structure

    @property
    def is_magnetic(self) -> bool:
        """Convenience property, returns True if any non-zero magmoms present.
        """
        return any(map(abs, self.structure.site_properties["magmom"]))

    @property
    def magmoms(self) -> np.ndarray:
        """Convenience property, returns magmoms as a numpy array.
        """

        return np.array(self.structure.site_properties["magmom"])

    @property
    def types_of_magnetic_species(self) -> Tuple[Union[Element, Species, DummySpecies], ...]:
        """Equivalent to Structure.types_of_specie but only returns
        magnetic species.

        Returns: types of Species as a list

        """
        if self.number_of_magnetic_sites > 0:
            structure = self.get_structure_with_only_magnetic_atoms()
            return tuple(sorted(structure.types_of_species))
        else:
            return tuple()

    @property
    def types_of_magnetic_specie(self) -> Tuple[Union[Element, Species, DummySpecies], ...]:
        """
        Specie->Species rename. Used to maintain backwards compatibility.
        """
        return self.types_of_magnetic_species

    @property
    def magnetic_species_and_magmoms(self) -> Dict[str, Any]:
        """Returns a dict of magnetic species and the magnitude of
        their associated magmoms. Will return a list if there are
        multiple magmoms per species.

        Returns: dict of magnetic species and magmoms
        """

        structure = self.get_ferromagnetic_structure()

        magtypes: Dict = {
            str(site.specie): set()
            for site in structure
            if site.properties["magmom"] != 0
        }

        for site in structure:
            if site.properties["magmom"] != 0:
                magtypes[str(site.specie)].add(site.properties["magmom"])

        for sp, magmoms in magtypes.items():
            if len(magmoms) == 1:
                magtypes[sp] = magmoms.pop()
            else:
                magtypes[sp] = sorted(list(magmoms))

        return magtypes

    @property
    def number_of_magnetic_sites(self) -> int:
        """Number of magnetic sites present in structure."""
        return int(np.sum([abs(m) > 0 for m in self.magmoms]))

    def number_of_unique_magnetic_sites(
        self, symprec: float = 1e-3, angle_tolerance: float = 5
    ) -> int:
        """

        Args:
          symprec: same as in SpacegroupAnalyzer (Default value = 1e-3)
          angle_tolerance: same as in SpacegroupAnalyzer (Default value = 5)

        Returns: Number of symmetrically-distinct magnetic sites present
        in structure.

        """

        structure = self.get_nonmagnetic_structure()

        sga = SpacegroupAnalyzer(
            structure, symprec=symprec, angle_tolerance=angle_tolerance
        )

        symm_structure = sga.get_symmetrized_structure()

        num_unique_mag_sites = 0

        for group_of_sites in symm_structure.equivalent_sites:
            if group_of_sites[0].specie in self.types_of_magnetic_species:
                num_unique_mag_sites += 1

        return num_unique_mag_sites

    @property
    def ordering(self) -> Ordering:
        """Applies heuristics to return a magnetic ordering for a collinear
        magnetic structure. Result is not guaranteed for correctness.

        Returns: Ordering Enum ('FiM' is used as the abbreviation for
        ferrimagnetic)
        """

        if not self.is_collinear:
            warnings.warn(
                "Detecting ordering in non-collinear structures not yet implemented."
            )
            return Ordering.Unknown

        if "magmom" not in self.structure.site_properties:
            # maybe this was a non-spin-polarized calculation, or we've
            # lost the magnetic moment information
            return Ordering.Unknown

        magmoms = self.magmoms

        max_magmom = max(magmoms)

        total_magnetization = abs(sum(magmoms))

        is_potentially_ferromagnetic = np.all(magmoms >= 0) or np.all(magmoms <= 0)

        if total_magnetization > 0 and is_potentially_ferromagnetic:
            return Ordering.FM
        elif total_magnetization > 0:
            return Ordering.FiM
        elif max_magmom > 0:
            return Ordering.AFM
        else:
            return Ordering.NM

    def get_exchange_group_info(
        self, symprec: float = 1e-2, angle_tolerance: float = 5.0
    ) -> Tuple[str, int]:
        """Returns the information on the symmetry of the Hamiltonian
        describing the exchange energy of the system, taking into
        account relative direction of magnetic moments but not their
        absolute direction.

        This is not strictly accurate (e.g. some/many atoms will
        have zero magnetic moments), but defining symmetry this
        way is a useful way of keeping track of distinct magnetic
        orderings within pymatgen.

        Args:
          symprec: same as SpacegroupAnalyzer (Default value = 1e-2)
          angle_tolerance: same as SpacegroupAnalyzer (Default value = 5.0)

        Returns:
          spacegroup_symbol, international_number

        """

        structure = self.get_structure_with_spin()

        return structure.get_space_group_info(
            symprec=symprec, angle_tolerance=angle_tolerance
        )

    def matches_ordering(self, other: Structure) -> bool:
        """Compares the magnetic orderings of one structure with another.

        Args:
          other: Structure to compare

        Returns: True or False
        """

        a = CollinearMagneticStructureAnalyzer(
            self.structure, overwrite_magmom_mode="normalize"
        ).get_structure_with_spin()

        # sign of spins doesn't matter, so we're comparing both
        # positive and negative versions of the structure
        # this code is possibly redundant, but is included out of
        # an abundance of caution
        b_positive = CollinearMagneticStructureAnalyzer(
            other, overwrite_magmom_mode="normalize", make_primitive=False
        )

        b_negative = b_positive.structure.copy()
        b_negative.add_site_property(
            "magmom", np.multiply(-1, b_negative.site_properties["magmom"])
        )

        b_negative = CollinearMagneticStructureAnalyzer(
            b_negative, overwrite_magmom_mode="normalize", make_primitive=False
        )

        b_positive = b_positive.get_structure_with_spin()
        b_negative = b_negative.get_structure_with_spin()

        if a.matches(b_positive) or a.matches(
            b_negative
        ):  # sometimes returns None (bug?)
            return True
        else:
            return False

    def __str__(self):
        """
        Sorts a Structure (by fractional co-ordinate), and
        prints sites with magnetic information. This is
        useful over Structure.__str__ because sites are in
        a consistent order, which makes visual comparison between
        two identical Structures with different magnetic orderings
        easier.
        """

        frac_coords = self.structure.frac_coords
        sorted_indices = np.lexsort(
            (frac_coords[:, 2], frac_coords[:, 1], frac_coords[:, 0])
        )
        s = Structure.from_sites([self.structure[idx] for idx in sorted_indices])

        # adapted from Structure.__repr__
        outs = ["Structure Summary", repr(s.lattice)]
        outs.append("Magmoms Sites")
        for site in s:
            if site.properties["magmom"] != 0:
                prefix = "{:+.2f}   ".format(site.properties["magmom"])
            else:
                prefix = "        "
            outs.append(prefix + repr(site))
        return "\n".join(outs)

class MagneticStructureEnumerator:
    """Combines MagneticStructureAnalyzer and MagOrderingTransformation to
    automatically generate a set of transformations for a given structure
    and produce a list of plausible magnetic orderings.
    """

    available_strategies = (
        "ferromagnetic",
        "antiferromagnetic",
        "ferrimagnetic_by_motif",
        "ferrimagnetic_by_species",
        "antiferromagnetic_by_motif",
        "nonmagnetic",
    )

    def __init__(
        self,
        structure: Structure,
        default_magmoms: Optional[Dict[str, float]] = None,
        strategies: Union[List[str], Tuple[str, ...]] = ("ferromagnetic", "antiferromagnetic"),
        automatic: bool = True,
        truncate_by_symmetry: bool = True,
        transformation_kwargs: Optional[Dict] = None,
    ):
        """
        This class will try generated different collinear
        magnetic orderings for a given input structure.

        If the input structure has magnetic moments defined, it
        is possible to use these as a hint as to which elements are
        magnetic, otherwise magnetic elements will be guessed
        (this can be changed using default_magmoms kwarg).

        Args:
            structure: input structure
            default_magmoms: (optional, defaults provided) dict of
                magnetic elements to their initial magnetic moments in µB, generally
                these are chosen to be high-spin since they can relax to a low-spin
                configuration during a DFT electronic configuration
            strategies: different ordering strategies to use, choose from:
                ferromagnetic, antiferromagnetic, antiferromagnetic_by_motif,
                ferrimagnetic_by_motif and ferrimagnetic_by_species (here, "motif",
                means to use a different ordering parameter for symmetry inequivalent
                sites)
            automatic: if True, will automatically choose sensible strategies
            truncate_by_symmetry: if True, will remove very unsymmetrical
                orderings that are likely physically implausible
            transformation_kwargs: keyword arguments to pass to
                MagOrderingTransformation, to change automatic cell size limits, etc.
        """

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

        self.structure = structure

        # decides how to process input structure, which sites are magnetic
        self.default_magmoms = default_magmoms

        # different strategies to attempt, default is usually reasonable
        self.strategies = list(strategies)
        # and whether to automatically add strategies that may be appropriate
        self.automatic = automatic

        # and whether to discard low symmetry structures
        self.truncate_by_symmetry = truncate_by_symmetry

        # other settings
        self.num_orderings = 64
        self.max_unique_sites = 8

        # kwargs to pass to transformation (ultimately to enumlib)
        default_transformation_kwargs = {"check_ordered_symmetry": False, "timeout": 5}
        transformation_kwargs = transformation_kwargs or {}
        transformation_kwargs.update(default_transformation_kwargs)
        self.transformation_kwargs = transformation_kwargs

        # our magnetically ordered structures will be
        # stored here once generated and also store which
        # transformation created them, this is used for
        # book-keeping/user interest, and
        # is be a list of strings in ("fm", "afm",
        # "ferrimagnetic_by_species", "ferrimagnetic_by_motif",
        # "afm_by_motif", "input_structure")
        self.ordered_structures: List[Structure] = []
        self.ordered_structure_origins: List[str] = []

        formula = structure.composition.reduced_formula

        # to process disordered magnetic structures, first make an
        # ordered approximation
        if not structure.is_ordered:
            raise ValueError(
                "Please obtain an ordered approximation of the "
                "input structure ({}).".format(formula)
            )

        # CollinearMagneticStructureAnalyzer is used throughout:
        # it can tell us whether the input is itself collinear (if not,
        # this workflow is not appropriate), and has many convenience
        # methods e.g. magnetic structure matching, etc.
        self.input_analyzer = CollinearMagneticStructureAnalyzer(
            structure, default_magmoms=default_magmoms, overwrite_magmom_mode="none"
        )

        # this workflow enumerates structures with different combinations
        # of up and down spin and does not include spin-orbit coupling:
        # if your input structure has vector magnetic moments, this
        # workflow is not appropriate
        if not self.input_analyzer.is_collinear:
            raise ValueError("Input structure ({}) is non-collinear.".format(formula))

        self.sanitized_structure = self._sanitize_input_structure(structure)

        # we will first create a set of transformations
        # and then apply them to our input structure
        self.transformations = self._generate_transformations(self.sanitized_structure)
        self._generate_ordered_structures(
            self.sanitized_structure, self.transformations
        )

    @staticmethod
    def _sanitize_input_structure(input_structure: Structure) -> Structure:
        """Sanitize our input structure by removing magnetic information
        and making primitive.

        Args:
          input_structure: Structure

        Returns: Structure

        """

        input_structure = input_structure.copy()

        # remove any annotated spin
        input_structure.remove_spin()

        # sanitize input structure: first make primitive ...
        input_structure = input_structure.get_primitive_structure(use_site_props=False)

        # ... and strip out existing magmoms, which can cause conflicts
        # with later transformations otherwise since sites would end up
        # with both magmom site properties and Species spins defined
        if "magmom" in input_structure.site_properties:
            input_structure.remove_site_property("magmom")

        return input_structure

    def _generate_transformations(
        self, structure: Structure
    ) -> Dict[str, MagOrderingTransformation]:
        """The central problem with trying to enumerate magnetic orderings is
        that we have to enumerate orderings that might plausibly be magnetic
        ground states, while not enumerating orderings that are physically
        implausible. The problem is that it is not always obvious by e.g.
        symmetry arguments alone which orderings to prefer. Here, we use a
        variety of strategies (heuristics) to enumerate plausible orderings,
        and later discard any duplicates that might be found by multiple
        strategies. This approach is not ideal, but has been found to be
        relatively robust over a wide range of magnetic structures.

        Args:
          structure: A sanitized input structure (_sanitize_input_structure)
        Returns: A dict of a transformation class instance (values) and name of
        enumeration strategy (keys)

        Returns: dict of Transformations keyed by strategy

        """

        formula = structure.composition.reduced_formula
        transformations: Dict[str, MagOrderingTransformation] = {}

        # analyzer is used to obtain information on sanitized input
        analyzer = CollinearMagneticStructureAnalyzer(
            structure,
            default_magmoms=self.default_magmoms,
            overwrite_magmom_mode="replace_all",
        )

        if not analyzer.is_magnetic:
            raise ValueError(
                "Not detected as magnetic, add a new default magmom for the "
                "element you believe may be magnetic?"
            )

        # now we can begin to generate our magnetic orderings
        self.logger.info("Generating magnetic orderings for {}".format(formula))

        mag_species_spin = analyzer.magnetic_species_and_magmoms
        types_mag_species = sorted(
            analyzer.types_of_magnetic_species,
            key=lambda sp: analyzer.default_magmoms.get(str(sp), 0),
            reverse=True,
        )
        num_mag_sites = analyzer.number_of_magnetic_sites
        num_unique_sites = analyzer.number_of_unique_magnetic_sites()

        # enumerations become too slow as number of unique sites (and thus
        # permutations) increase, 8 is a soft limit, this can be increased
        # but do so with care
        if num_unique_sites > self.max_unique_sites:
            raise ValueError("Too many magnetic sites to sensibly perform enumeration.")

        # maximum cell size to consider: as a rule of thumb, if the primitive cell
        # contains a large number of magnetic sites, perhaps we only need to enumerate
        # within one cell, whereas on the other extreme if the primitive cell only
        # contains a single magnetic site, we have to create larger supercells
        if "max_cell_size" not in self.transformation_kwargs:
            # TODO: change to 8 / num_mag_sites ?
            self.transformation_kwargs["max_cell_size"] = max(1, int(4 / num_mag_sites))
        self.logger.info(
            "Max cell size set to {}".format(
                self.transformation_kwargs["max_cell_size"]
            )
        )

        # when enumerating ferrimagnetic structures, it's useful to detect
        # symmetrically distinct magnetic sites, since different
        # local environments can result in different magnetic order
        # (e.g. inverse spinels)
        # initially, this was done by co-ordination number, but is
        # now done by a full symmetry analysis
        sga = SpacegroupAnalyzer(structure)
        structure_sym = sga.get_symmetrized_structure()
        wyckoff = ["n/a"] * len(structure)
        for indices, symbol in zip(
            structure_sym.equivalent_indices, structure_sym.wyckoff_symbols
        ):
            for index in indices:
                wyckoff[index] = symbol
        is_magnetic_sites = [
            True if site.specie in types_mag_species else False for site in structure
        ]
        # we're not interested in sites that we don't think are magnetic,
        # set these symbols to None to filter them out later
        wyckoff = [
            symbol if is_magnetic_site else "n/a"
            for symbol, is_magnetic_site in zip(wyckoff, is_magnetic_sites)
        ]
        structure.add_site_property("wyckoff", wyckoff)
        wyckoff_symbols = set(wyckoff) - {"n/a"}

        # if user doesn't specifically request ferrimagnetic orderings,
        # we apply a heuristic as to whether to attempt them or not
        if self.automatic:
            if (
                "ferrimagnetic_by_motif" not in self.strategies
                and len(wyckoff_symbols) > 1
                and len(types_mag_species) == 1
            ):
                self.strategies += ["ferrimagnetic_by_motif"]

            if (
                "antiferromagnetic_by_motif" not in self.strategies
                and len(wyckoff_symbols) > 1
                and len(types_mag_species) == 1
            ):
                self.strategies += ["antiferromagnetic_by_motif"]

            if (
                "ferrimagnetic_by_species" not in self.strategies
                and len(types_mag_species) > 1
            ):
                self.strategies += ["ferrimagnetic_by_species"]

        # we start with a ferromagnetic ordering
        if "ferromagnetic" in self.strategies:
            # TODO: remove 0 spins !

            fm_structure = analyzer.get_ferromagnetic_structure()
            # store magmom as spin property, to be consistent with output from
            # other transformations
            fm_structure.add_spin_by_site(fm_structure.site_properties["magmom"])
            fm_structure.remove_site_property("magmom")

            # we now have our first magnetic ordering...
            self.ordered_structures.append(fm_structure)
            self.ordered_structure_origins.append("fm")

        # we store constraint(s) for each strategy first,
        # and then use each to perform a transformation later
        all_constraints: Dict[str, Any] = {}

        # ...to which we can add simple AFM cases first...
        if "antiferromagnetic" in self.strategies:

            constraint = MagOrderParameterConstraint(
                0.5,
                # TODO: update MagOrderParameterConstraint in
                # pymatgen to take types_mag_species directly
                species_constraints=list(map(str, types_mag_species)),
            )
            all_constraints["afm"] = [constraint]

            # allows for non-magnetic sublattices
            if len(types_mag_species) > 1:
                for sp in types_mag_species:
                    constraints = [
                        MagOrderParameterConstraint(0.5, species_constraints=str(sp))
                    ]

                    all_constraints["afm_by_{}".format(sp)] = constraints

        # ...and then we also try ferrimagnetic orderings by motif if a
        # single magnetic species is present...
        if "ferrimagnetic_by_motif" in self.strategies and len(wyckoff_symbols) > 1:

            # these orderings are AFM on one local environment, and FM on the rest
            for symbol in wyckoff_symbols:
                constraints = [
                    MagOrderParameterConstraint(
                        0.5, site_constraint_name="wyckoff", site_constraints=symbol
                    ),
                    MagOrderParameterConstraint(
                        1.0,
                        site_constraint_name="wyckoff",
                        site_constraints=list(wyckoff_symbols - {symbol}),
                    ),
                ]

                all_constraints["ferri_by_motif_{}".format(symbol)] = constraints

        # and also try ferrimagnetic when there are multiple magnetic species
        if "ferrimagnetic_by_species" in self.strategies:

            sp_list = [str(site.specie) for site in structure]
            num_sp = {sp: sp_list.count(str(sp)) for sp in types_mag_species}
            total_mag_sites = sum(num_sp.values())

            for sp in types_mag_species:
                # attempt via a global order parameter
                all_constraints["ferri_by_{}".format(sp)] = num_sp[sp] / total_mag_sites

                # attempt via afm on sp, fm on remaining species

                constraints = [
                    MagOrderParameterConstraint(0.5, species_constraints=str(sp)),
                    MagOrderParameterConstraint(
                        1.0,
                        species_constraints=list(
                            map(str, set(types_mag_species) - {sp})
                        ),
                    ),
                ]

                all_constraints["ferri_by_{}_afm".format(sp)] = constraints

        # ...and finally, we can try orderings that are AFM on one local
        # environment, and non-magnetic on the rest -- this is less common
        # but unless explicitly attempted, these states are unlikely to be found
        if "antiferromagnetic_by_motif" in self.strategies:

            for symbol in wyckoff_symbols:
                constraints = [
                    MagOrderParameterConstraint(
                        0.5, site_constraint_name="wyckoff", site_constraints=symbol
                    )
                ]

                all_constraints["afm_by_motif_{}".format(symbol)] = constraints

        # and now construct all our transformations for each strategy
        transformations = {}
        for name, constraints in all_constraints.items():
            trans = MagOrderingTransformation(
                mag_species_spin,
                order_parameter=constraints,
                **self.transformation_kwargs
            )

            transformations[name] = trans

        return transformations

    def _generate_ordered_structures(
        self,
        sanitized_input_structure: Structure,
        transformations: Dict[str, MagOrderingTransformation],
    ):
        """Apply our input structure to our list of transformations and output a list
        of ordered structures that have been pruned for duplicates and for those
        with low symmetry (optional).

        Args:
            sanitized_input_structure: A sanitized input structure
            (_sanitize_input_structure)
            transformations: A dict of transformations (values) and name of
            enumeration strategy (key), the enumeration strategy name is just
            for record keeping

        Returns: None (sets self.ordered_structures
        and self.ordered_structures_origins instance variables)

        Returns: List of Structures
        """

        ordered_structures = self.ordered_structures
        ordered_structures_origins = self.ordered_structure_origins

        # utility function to combine outputs from several transformations
        def _add_structures(
            ordered_structures, ordered_structures_origins, structures_to_add, origin=""
        ):
            """Transformations with return_ranked_list can return either
            just Structures or dicts (or sometimes lists!) -- until this
            is fixed, we use this function to concat structures given
            by the transformation.
            """
            if structures_to_add:
                # type conversion
                if isinstance(structures_to_add, Structure):
                    structures_to_add = [structures_to_add]
                structures_to_add = [
                    s["structure"] if isinstance(s, dict) else s
                    for s in structures_to_add
                ]
                # concatenation
                ordered_structures += structures_to_add
                ordered_structures_origins += [origin] * len(structures_to_add)
                self.logger.info(
                    "Adding {} ordered structures: {}".format(
                        len(structures_to_add), origin
                    )
                )

            return ordered_structures, ordered_structures_origins

        for origin, trans in self.transformations.items():
            structures_to_add = trans.apply_transformation(
                self.sanitized_structure, return_ranked_list=self.num_orderings
            )
            ordered_structures, ordered_structures_origins = _add_structures(
                ordered_structures,
                ordered_structures_origins,
                structures_to_add,
                origin=origin,
            )

        # in case we've introduced duplicates, let's remove them
        self.logger.info("Pruning duplicate structures.")
        structures_to_remove: List[int] = []
        for idx, ordered_structure in enumerate(ordered_structures):
            if idx not in structures_to_remove:
                duplicate_checker = CollinearMagneticStructureAnalyzer(
                    ordered_structure, overwrite_magmom_mode="none"
                )
                for check_idx, check_structure in enumerate(ordered_structures):
                    if check_idx not in structures_to_remove and check_idx != idx:
                        if duplicate_checker.matches_ordering(check_structure):
                            structures_to_remove.append(check_idx)

        if len(structures_to_remove):
            self.logger.info(
                "Removing {} duplicate ordered structures".format(
                    len(structures_to_remove)
                )
            )
            ordered_structures = [
                s
                for idx, s in enumerate(ordered_structures)
                if idx not in structures_to_remove
            ]
            ordered_structures_origins = [
                o
                for idx, o in enumerate(ordered_structures_origins)
                if idx not in structures_to_remove
            ]

        # also remove low symmetry structures
        if self.truncate_by_symmetry:

            # by default, keep structures with 5 most symmetric space groups
            if not isinstance(self.truncate_by_symmetry, int):
                self.truncate_by_symmetry = 5

            self.logger.info("Pruning low symmetry structures.")

            # first get a list of symmetries present
            symmetry_int_numbers = [
                s.get_space_group_info()[1] for s in ordered_structures
            ]

            # then count the number of symmetry operations for that space group
            num_sym_ops = [
                len(SpaceGroup.from_int_number(n).symmetry_ops)
                for n in symmetry_int_numbers
            ]

            # find the largest values...
            max_symmetries = sorted(list(set(num_sym_ops)), reverse=True)

            # ...and decide which ones to keep
            if len(max_symmetries) > self.truncate_by_symmetry:
                max_symmetries = max_symmetries[0:5]
            structs_to_keep = [
                (idx, num)
                for idx, num in enumerate(num_sym_ops)
                if num in max_symmetries
            ]

            # sort so that highest symmetry structs are first
            structs_to_keep = sorted(
                structs_to_keep, key=lambda x: (x[1], -x[0]), reverse=True
            )

            self.logger.info(
                "Removing {} low symmetry "
                "ordered structures".format(
                    len(ordered_structures) - len(structs_to_keep)
                )
            )

            ordered_structures = [ordered_structures[i] for i, _ in structs_to_keep]
            ordered_structures_origins = [
                ordered_structures_origins[i] for i, _ in structs_to_keep
            ]

            # and ensure fm is always at index 0
            fm_index = ordered_structures_origins.index("fm")
            ordered_structures.insert(0, ordered_structures.pop(fm_index))
            ordered_structures_origins.insert(
                0, ordered_structures_origins.pop(fm_index)
            )

        # if our input structure isn't in our generated structures,
        # let's add it manually and also keep a note of which structure
        # is our input: this is mostly for book-keeping/benchmarking
        self.input_index = None
        self.input_origin = None
        if self.input_analyzer.ordering != Ordering.NM:
            matches = [
                self.input_analyzer.matches_ordering(s) for s in ordered_structures
            ]
            if not any(matches):
                ordered_structures.append(self.input_analyzer.structure)
                ordered_structures_origins.append("input")
                self.logger.info(
                    "Input structure not present in enumerated structures, adding..."
                )
            else:
                self.logger.info(
                    "Input structure was found in enumerated "
                    "structures at index {}".format(matches.index(True))
                )
                self.input_index = matches.index(True)
                self.input_origin = ordered_structures_origins[self.input_index]

        self.ordered_structures = ordered_structures
        self.ordered_structure_origins = ordered_structures_origins

MagneticDeformation = namedtuple("MagneticDeformation", "type deformation")

def magnetic_deformation(
    structure_A: Structure, structure_B: Structure
) -> MagneticDeformation:
    """Calculates 'magnetic deformation proxy',
    a measure of deformation (norm of finite strain)
    between 'non-magnetic' (non-spin-polarized) and
    ferromagnetic structures.

    Adapted from Bocarsly et al. 2017,
    doi: 10.1021/acs.chemmater.6b04729

    Args:
      structure_A: Structure
      structure_B: Structure

    Returns: Magnetic deformation
    """

    # retrieve orderings of both input structures
    ordering_a = CollinearMagneticStructureAnalyzer(
        structure_A, overwrite_magmom_mode="none"
    ).ordering
    ordering_b = CollinearMagneticStructureAnalyzer(
        structure_B, overwrite_magmom_mode="none"
    ).ordering

    # get a type string, this is either 'NM-FM' for between non-magnetic
    # and ferromagnetic, as in Bocarsly paper, or e.g. 'FM-AFM'
    type_str = "{}-{}".format(ordering_a.value, ordering_b.value)

    lattice_a = structure_A.lattice.matrix.T
    lattice_b = structure_B.lattice.matrix.T
    lattice_a_inv = np.linalg.inv(lattice_a)
    p = np.dot(lattice_a_inv, lattice_b)
    eta = 0.5 * (np.dot(p.T, p) - np.identity(3))
    w, v = np.linalg.eig(eta)
    deformation = 100 * (1.0 / 3.0) * np.sqrt(w[0] ** 2 + w[1] ** 2 + w[2] ** 2)

    return MagneticDeformation(deformation=deformation, type=type_str)


import re
import sys
import warnings

import numpy as np
from scipy.optimize import fsolve

# suppress sqrt warning. it's going to happen. Don't really care if it does.
warnings.filterwarnings('ignore', 'invalid value encountered in sqrt')

def new_duct_network():
    ducts = dict(title=None, fan_pressure=None, air_density=None, roughness=None, rounding=None, fittings=[])
    return ducts

def new_fitting():
    fitting = dict(ID=None, type=None, IDup=None, BranchUP=None, IDdownMain=None, IDdownBranch=None, flow=None,
                   flowMain=None, flowBranch=None, size=None, sizeMain=None, sizeBranch=None,
                   pdrop=None, pdropMain=None, pdropBranch=None, length=None, fandist=None, diffuser_psum=None)
    return fitting

def read_input_file(filename):
    file = open(filename, 'r')
    data = file.readlines()
    file.close()
    return data

def process_keywords(data):
    ducts = new_duct_network()
    for line in data:
        line = line.lower()
        if line.find('#') != -1:
            continue  # comment line. not the lines we are looking for.
        else:
            item = [x.strip() for x in line.split(',')]
            if item[0] == 'title':
                ducts['title'] = item[1]
            elif item[0] == 'fan_pressure':
                ducts['fan_pressure'] = float(item[1])
            elif item[0] == 'air_density':
                ducts['air_density'] = float(item[1])
            elif item[0] == 'roughness':
                ducts['roughness'] = float(item[1])
            elif item[0] == 'rounding':
                ducts['rounding'] = item[1]
            elif item[0] == 'fitting':  # initializing fittings information
                fitting = new_fitting()
                fitting['fandist'] = 0
                fitting['ID'] = float(item[1])
                fitting['type'] = item[2]
                try:  # checking for air handeling units
                    fitting['IDup'] = item[3]
                except Exception:
                    pass
                if item[2] == 'duct':  # check for ducts. Ducts should have length feet
                    fitting['length'] = float(item[4])
                elif item[2] == 'diffuser':  # check for diffusers. Diffusers should have flowrate in CFM
                    fitting['flow'] = float(item[4])
                else:  # Everything not a duct or diffuser
                    pass
                ducts['fittings'].append(fitting)
            else:
                continue
    return ducts

def find_fitting(ID, fittings):
    for fitting in fittings:
        if fitting['ID'] == ID:
            return fitting

def make_connections(fittings):
    # regex patterns for matching main and branch ID up text
    main_pattern = re.compile(r'\d+\b-main\b')
    branch_pattern = re.compile(r'\d+\b-branch\b')
    for fitting in fittings:
        if fitting['IDup'] is None:  # check for AHUs
            continue
        elif main_pattern.match(fitting['IDup']):  # main line from Tee
            # Find index of hyphen when present
            index_of_hyphen = fitting['IDup'].find('-')
            tee_ID = int(fitting['IDup'][:index_of_hyphen])
            tee_fitting = find_fitting(tee_ID, fittings)
            tee_fitting['IDdownMain'] = fitting['ID']
        elif branch_pattern.match(fitting['IDup']):  # branch from Tee
            index_of_hyphen = fitting['IDup'].find('-')
            tee_ID = int(fitting['IDup'][:index_of_hyphen])
            tee_fitting = find_fitting(tee_ID, fittings)
            tee_fitting['IDdownBranch'] = fitting['ID']
        else:  # all other fittings
            fittingUp = find_fitting(float(fitting['IDup']), fittings)
            fittingUp['IDdownMain'] = fitting['ID']

def setup_fan_distances(fittings):
    for fitting in fittings:
        # check for AHU
        if fitting['type'] == 'air_handling_unit':
            fitting['fandist'] = 0
        # fittingUp is a tee
        elif fitting['IDup'].find('-') != -1:
            index_of_hyphen = fitting['IDup'].find('-')
            tee_ID = int(fitting['IDup'][:index_of_hyphen])
            fittingUp = find_fitting(tee_ID, fittings)
            fitting['fandist'] = fittingUp['fandist']
        # fittings up is duct. Add length
        elif find_fitting(int(fitting['IDup']), fittings)['type'] == 'duct':
            fittingUp = find_fitting(int(fitting['IDup']), fittings)
            fitting['fandist'] = fittingUp['fandist'] + fittingUp['length']
        # all the elbow, tee, and duct fittings
        else:
            fittingUp = find_fitting(int(fitting['IDup']), fittings)
            fitting['fandist'] = fittingUp['fandist']

def setup_flowrates(fittings):  # Iterates, takes flow from duct downstream and makes it the flow of the fitting
    for i in range(1000):  # Iterations
        for fitting in fittings:
            if fitting['IDdownBranch'] is None and fitting['IDdownMain'] is not None:  # Straightaways
                MainDown = find_fitting(fitting['IDdownMain'], fittings)
                flow = MainDown['flow']
                if flow is None:  # Avoids errors
                    continue
                fitting['flow'] = flow  # Sets flowrate to singular downstream piece
            elif fitting['IDdownMain'] is None:  # Diffusers
                continue
            elif fitting['IDdownBranch'] is not None and fitting['IDdownMain'] is not None:  # Tee's
                MainDown = find_fitting(fitting['IDdownMain'], fittings)
                flow = MainDown['flow']  # Sets partial flow to main downstream piece
                if flow is None:  # Avoids errors
                    continue
                BranchDown = find_fitting(fitting['IDdownBranch'], fittings)
                if BranchDown['flow'] is None:  # Avoids errors
                    continue
                flow = BranchDown['flow'] + flow  # Adds branch flow to mainflow
                fitting['flow'] = flow  # Sets flow for fitting
            else:
                print("There was an error in the flow rate.")  # Error message, just in case

def get_little_f(dia, velocity, roughness):
    def func(vals):
        f = vals
        Re = 8.5 * (dia / 12) * velocity  # eqn (21)

        # eqn (19) 2013 version corrected for units
        left_side = 1 / np.sqrt(f)
        right_right = (-1) * 2 * np.log10((roughness / (3.7 * (dia / 12))) + (2.51 / (Re * np.sqrt(f))))
        return (right_right - left_side)

    guess = 10
    finished = False
    while not finished:
        f = fsolve(func, guess, full_output=True)
        if int(f[2]) == 1:  # check solution flag
            finished = True
        else:
            guess = guess / 2.0  # update guess
    return f[0][0]

def largest_path(fittings):  # Finds the diffuser with the longest path to the fan
    fitting_compare = find_fitting(1, fittings)  # Sets the initial fitting that will be compared
    for fitting in fittings:
        if fitting['type'] == 'diffuser':
            if fitting['fandist'] > fitting_compare['fandist']:  # compare fan distances
                fitting_compare = fitting  # sets new comparison
    return fitting_compare

def duct_pressure_drop(dia, flow, length, density, roughness):  # dia [inches], length [ft]
    area = (np.pi * (dia / 12) ** 2) / 4  # [ft^2]
    velocity = flow / area  # [ft/min=fpm]
    f = get_little_f(dia, velocity, roughness)
    pdrop = ((12 * f * length) / (dia / 12)) * density * (velocity / 1097) ** 2
    return pdrop

# Nick and Charlie
def pressure_drop_sum(ID, fittings):  # calculates total pressure loss of ANY RUN
    # find the fitting dictionary attached to the ID in quesiton
    fitting = find_fitting(int(ID), fittings)
    route = [int(fitting['ID'])]
    main_pattern = re.compile(r'\d+\b-main\b')
    branch_pattern = re.compile(r'\d+\b-branch\b')

    # find route IDs
    while fitting['type'] != 'air_handling_unit':  # while not at the air handler
        if main_pattern.match(fitting['IDup']):  # matching main if IDup is tee
            index_of_hyphen = fitting['IDup'].find('-')
            ID = int(fitting['IDup'][:index_of_hyphen])
            route.append(ID)
            fitting = find_fitting(ID, fittings)
        elif branch_pattern.match(fitting['IDup']):  # matching branch if IDup is tee
            index_of_hyphen = fitting['IDup'].find('-')
            ID = int(fitting['IDup'][:index_of_hyphen])
            route.append(ID)
            fitting = find_fitting(ID, fittings)
        else:
            ID = int(fitting['IDup'])
            route.append(ID)
            fitting = find_fitting(ID, fittings)

    pdrop_sum = 0
    for i in range(len(route)):
        fitting = find_fitting(int(route[i]), fittings)
        if fitting['type'] == 'duct' or fitting['type'] == 'elbow':
            pdrop_sum += fitting['pdrop']
        elif fitting['type'] == 'tee':
            # next downstream fitting ID
            next_fitting_ID = route[i - 1]
            if next_fitting_ID == fitting['IDdownMain']:
                pdrop_sum += find_fitting(int(next_fitting_ID), fittings)['pdropMain']
            elif next_fitting_ID == fitting['IDdownBranch']:
                pdrop_sum += find_fitting(int(next_fitting_ID), fittings)['pdropBranch']
    return pdrop_sum

def fitting_loss_sum(fittings):  # calculates total pressure loss of tees and elbows only, of LONGEST RUN
    farthest_fitting = largest_path(fittings)
    fitting = farthest_fitting
    longest_route = [int(fitting['ID'])]
    main_pattern = re.compile(r'\d+\b-main\b')
    branch_pattern = re.compile(r'\d+\b-branch\b')

    # find longest route using farthest diffusers
    while fitting['type'] != 'air_handling_unit':  # while not at the air handler
        if main_pattern.match(fitting['IDup']):  # matching main if IDup is tee
            index_of_hyphen = fitting['IDup'].find('-')
            ID = int(fitting['IDup'][:index_of_hyphen])
            longest_route.append(ID)
            fitting = find_fitting(ID, fittings)
        elif branch_pattern.match(fitting['IDup']):  # matching branch if IDup is tee
            index_of_hyphen = fitting['IDup'].find('-')
            ID = int(fitting['IDup'][:index_of_hyphen])
            longest_route.append(ID)
            fitting = find_fitting(ID, fittings)
        else:
            ID = int(fitting['IDup'])
            longest_route.append(ID)
            fitting = find_fitting(ID, fittings)

    pdrop_fitloss_sum = 0
    for i in range(len(longest_route)):
        fitting = find_fitting(int(longest_route[i]), fittings)
        if fitting['type'] == 'elbow':
            pdrop_fitloss_sum += fitting['pdrop']
        elif fitting['type'] == 'tee':
            # next downstream fitting ID
            next_fitting_ID = longest_route[i - 1]
            if next_fitting_ID == fitting['IDdownMain']:
                pdrop_fitloss_sum += find_fitting(int(next_fitting_ID), fittings)['pdropMain']
            elif next_fitting_ID == fitting['IDdownBranch']:
                pdrop_fitloss_sum += find_fitting(int(next_fitting_ID), fittings)['pdropBranch']
    return pdrop_fitloss_sum

def get_duct_size(deltap, flow, length, density, roughness):
    def func(vals):
        dia = vals
        return deltap - duct_pressure_drop(dia, flow, length, density, roughness)

    finished = False
    guess = .01
    while not finished:
        # print(guess)
        f = fsolve(func, guess, full_output=True)
        if int(f[2]) == 1:
            finished = True
        else:
            guess = guess * 2.0
    diameter = f[0][0]
    return diameter

def findBetween(x, xlist):
    # iterate over xlist to find where x fits in
    for position in range(0, len(xlist) - 1):
        # check if the current iteration in xlist is good
        if (xlist[position] <= x <= xlist[position + 1]) or (xlist[position] >= x >= xlist[position + 1]):
            return position

def interp1D(x, xlist, ylist):
    # find between value
    position = findBetween(x, xlist)

    # gather needed numbers from xlist and ylist
    x1 = xlist[position]
    x3 = xlist[position + 1]
    y1 = ylist[position]
    y3 = ylist[position + 1]

    # interpolate value of y
    y = (((x - x1) * (y3 - y1)) / (x3 - x1)) + y1
    return y

def interp2D(x, y, xlist, ylist, zmatrix):
    # find the correct column to use
    yposition = findBetween(y, ylist)

    # extract and interpolate each column
    x1 = interp1D(x, xlist, zmatrix[:, yposition])
    x2 = interp1D(x, xlist, zmatrix[:, yposition + 1])

    # interpolate between columns
    z = interp1D(y, ylist[yposition:yposition + 2], [x1, x2])
    return z

def tee_pressure_drop(dia, density, flow, outlet_flow, outlet_dia, branch):
    '''

    :param dia:
    :param density:
    :param flow:
    :param outlet_flow:
    :param outlet_dia:
    :param branch: Boolean
    :return:
    '''
    Q = np.array([.1, .2, .3, .4, .5, .6, .7, .8, .9])  # top most row Q_(branch/main)/Q_common
    A = np.array([.1, .2, .3, .4, .5, .6, .7, .8, .9])  # right most column A_(branch/main)/A_common

    # Table from 2009 ASHRAE Handbook 21.50 for SD5-10 Tee, Concial Branch Tapered into Body, Diverging
    # c_branch values
    sd5cb = np.array([[0.65, 0.24, 0.15, 0.11, 0.09, 0.07, 0.06, 0.05, 0.05],
                      [2.98, 0.65, 0.33, 0.24, 0.18, 0.15, 0.13, 0.11, 0.10],
                      [7.36, 1.56, 0.65, 0.39, 0.29, 0.24, 0.20, 0.17, 0.15],
                      [13.78, 2.98, 1.20, 0.65, 0.43, 0.33, 0.27, 0.24, 0.21],
                      [22.24, 4.92, 1.98, 1.04, 0.65, 0.47, 0.36, 0.30, 0.26],
                      [32.73, 7.36, 2.98, 1.56, 0.96, 0.65, 0.49, 0.39, 0.33],
                      [45.26, 10.32, 4.21, 2.21, 1.34, 0.90, 0.65, 0.51, 0.42],
                      [59.82, 13.78, 5.67, 2.98, 1.80, 1.20, 0.86, 0.65, 0.52],
                      [76.41, 17.75, 7.36, 3.88, 2.35, 1.56, 1.11, 0.83, 0.65]])
    # c_main values
    sd5cm = np.array([[0.13, 0.16, 0.57, 0.74, 0.74, 0.70, 0.65, 0.60, 0.56],
                      [0.20, 0.13, 0.15, 0.16, 0.28, 0.57, 0.69, 0.74, 0.75],
                      [0.90, 0.13, 0.13, 0.14, 0.15, 0.16, 0.20, 0.42, 0.57],
                      [2.88, 0.20, 0.14, 0.13, 0.14, 0.15, 0.15, 0.16, 0.34],
                      [6.25, 0.37, 0.17, 0.14, 0.13, 0.14, 0.14, 0.15, 0.15],
                      [11.88, 0.90, 0.20, 0.13, 0.14, 0.13, 0.14, 0.14, 0.15],
                      [18.62, 1.71, 0.33, 0.18, 0.16, 0.14, 0.13, 0.15, 0.14],
                      [26.88, 2.88, 0.50, 0.20, 0.15, 0.14, 0.13, 0.13, 0.14],
                      [36.45, 4.46, 0.90, 0.30, 0.19, 0.16, 0.15, 0.14, 0.13]])
    area = (np.pi * (dia / 12) ** 2) / 4
    velocity = flow / area
    p_v = density * (velocity / 1097) ** 2
    A_common = (np.pi * (dia / 12) ** 2) / 4
    A_outlet = (np.pi * (outlet_dia / 12) ** 2) / 4
    area_ratio = A_outlet / A_common
    flow_ratio = outlet_flow / flow

    # protect from extrapolating on the table
    if area_ratio > .9:
        area_ratio = .9
    elif area_ratio < .1:
        area_ratio = .1
    if flow_ratio > .9:
        flow_ratio = .9
    elif flow_ratio < .1:
        flow_ratio = .1
    if branch:
        c_branch = interp2D(area_ratio, flow_ratio, A, Q, sd5cb)
        pdrop = c_branch * p_v
    else:
        c_main = interp2D(area_ratio, flow_ratio, A, Q, sd5cm)
        pdrop = c_main * p_v
    return pdrop

def elbow_pressure_drop(dia, flow, density):
    # Table from ASHRAE 2009 chapter 21 for pleated 90 degree elbow
    D = np.array([4, 6, 8, 10, 12, 14, 16])
    Co = np.array([0.57, 0.43, 0.34, 0.28, 0.26, 0.25, 0.25])

    # protect from extrapolating on the table
    if dia >= 16:
        c_o = 0.25
    elif dia <= 4:
        c_o = 0.57
    else:
        c_o = interp1D(dia, D, Co)

    area = (np.pi * (dia / 12) ** 2) / 4
    velocity = flow / area
    p_v = density * (velocity / 1097) ** 2
    pdrop = c_o * p_v
    return pdrop

# Nick Nelsen 5/4/17
def sizing_iterate_nick(ducts):
    # assign for easier use later
    density = ducts['air_density']
    roughness = ducts['roughness']
    fan_pressure = ducts['fan_pressure']
    fittings = ducts['fittings']
    maxlength = largest_path(fittings)['fandist']

    # take care of nonetype errors the cheesy way
    for fitting in fittings:
        if fitting['pdrop'] is None:
            fitting['pdrop'] = 0.0
        if fitting['pdropMain'] is None:
            fitting['pdropMain'] = 0.0
        if fitting['pdropBranch'] is None:
            fitting['pdropBranch'] = 0.0

    psum = fitting_loss_sum(fittings)
    dpdl = (fan_pressure - psum) / maxlength
    dpdl_old = 0
    count = 0
    # main loop to size ducts first, then elbows, and finally tees
    while abs(dpdl - dpdl_old) >= 1e-10:
        for fitting in fittings:  # solving ducts
            if fitting['type'] == 'duct':
                length = fitting['length']
                deltap = dpdl * length
                flow = fitting['flow']
                diameter = get_duct_size(deltap, flow, length, density, roughness)
                fitting['size'] = diameter
                # p_duct = duct_pressure_drop(diameter, flow, length, density, roughness)
                p_duct = deltap
                # print(p_duct, p_duct2)
                fitting['pdrop'] = p_duct

        for fitting in fittings:  # solving elbows
            if fitting['type'] == 'elbow':
                down_fitting = find_fitting(int(fitting['IDdownMain']), fittings)
                if down_fitting['type'] == 'duct':
                    e_diameter = down_fitting['size']
                    p_elbow = elbow_pressure_drop(e_diameter, fitting['flow'], density)
                    fitting['size'] = e_diameter
                    fitting['pdrop'] = p_elbow

                # if up_fitting is a tee
                elif fitting['IDup'].find('-') != -1:
                    index_of_hyphen = fitting['IDup'].find('-')
                    tee_ID = int(fitting['IDup'][:index_of_hyphen])
                    up_fitting = find_fitting(tee_ID, fittings)
                    e_diameter = up_fitting['size']
                    p_elbow = elbow_pressure_drop(e_diameter, fitting['flow'], density)
                    fitting['size'] = e_diameter
                    fitting['pdrop'] = p_elbow
                # if up_fitting has no hyphen
                else:
                    up_fitting = find_fitting(int(fitting['IDup']), fittings)
                    # if up_fitting['size'] is not None:
                    e_diameter = up_fitting['size']
                    p_elbow = elbow_pressure_drop(e_diameter, fitting['flow'], density)
                    fitting['size'] = e_diameter
                    fitting['pdrop'] = p_elbow

        for fitting in fittings:  # solving tees
            if fitting['type'] == 'tee':
                up_fitting = find_fitting(int(fitting['IDup']), fittings)
                tee_inlet_diameter = up_fitting['size']
                p_tee_main = tee_pressure_drop(tee_inlet_diameter, density, fitting['flow'],
                                               find_fitting(fitting['IDdownMain'], fittings)['flow'],
                                               find_fitting(fitting['IDdownMain'], fittings)['size'], False)
                p_tee_branch = tee_pressure_drop(tee_inlet_diameter, density, fitting['flow'],
                                                 find_fitting(fitting['IDdownBranch'], fittings)['flow'],
                                                 find_fitting(fitting['IDdownBranch'], fittings)['size'], True)
                fitting['size'] = tee_inlet_diameter
                fitting['sizeMain'] = find_fitting(fitting['IDdownMain'], fittings)['size']
                fitting['sizeBranch'] = find_fitting(fitting['IDdownBranch'], fittings)['size']
                fitting['pdropMain'] = p_tee_main
                fitting['pdropBranch'] = p_tee_branch

        psum = fitting_loss_sum(fittings)
        dpdl_old = dpdl
        dpdl = (fan_pressure - psum) / maxlength

    for fitting in fittings:
        if fitting['type'] == 'diffuser':
            fitting['diffuser_psum'] = pressure_drop_sum(int(fitting['ID']), fittings)

    # rounding stuff
    if ducts['rounding'] is not None:
        if ducts['rounding'] == 'nearest':
            for fitting in fittings:
                if fitting['type'] != 'air_handling_unit' and fitting['type'] != 'diffuser':
                    fitting['size'] = round(fitting['size'])
                    if fitting['type'] == 'tee':
                        fitting['sizeMain'] = round(fitting['sizeMain'])
                        fitting['sizeBranch'] = round(fitting['sizeBranch'])
        elif ducts['rounding'] == 'up':
            for fitting in fittings:
                if fitting['type'] != 'air_handling_unit' and fitting['type'] != 'diffuser':
                    fitting['size'] = np.ceil(fitting['size'])
                    if fitting['type'] == 'tee':
                        fitting['sizeMain'] = np.ceil(fitting['sizeMain'])
                        fitting['sizeBranch'] = np.ceil(fitting['sizeBranch'])
        elif ducts['rounding'] == 'down':
            for fitting in fittings:
                if fitting['type'] != 'air_handling_unit' and fitting['type'] != 'diffuser':
                    fitting['size'] = np.floor(fitting['size'])
                    if fitting['type'] == 'tee':
                        fitting['sizeMain'] = np.floor(fitting['sizeMain'])
                        fitting['sizeBranch'] = np.floor(fitting['sizeBranch'])

        # recalculate pdrop everything after rounding
        for fitting in fittings:
            if fitting['type'] == 'duct':
                fitting['pdrop'] = duct_pressure_drop(fitting['size'], fitting['flow'], fitting['length'], density,
                                                      roughness)
            elif fitting['type'] == 'tee':
                fitting['flowMain'] = find_fitting(int(fitting['IDdownMain']), fittings)['flow']
                fitting['pdropMain'] = tee_pressure_drop(fitting['size'], density, fitting['flow'], fitting['flowMain'],
                                                         fitting['sizeMain'], False)
                fitting['flowBranch'] = find_fitting(int(fitting['IDdownBranch']), fittings)['flow']
                fitting['pdropBranch'] = tee_pressure_drop(fitting['size'], density, fitting['flow'],
                                                           fitting['flowBranch'], fitting['sizeBranch'], True)
            elif fitting['type'] == 'elbow':
                fitting['pdrop'] = elbow_pressure_drop(fitting['size'], fitting['flow'], density)

    # finally apply sizes to diffuers
    for fitting in fittings:
        if fitting['type'] == 'diffuser':
            fitting['size'] = find_fitting(int(fitting['IDup']), fittings)['size']

    return

def print_results(fittings):
    # file=open("pyductresults.txt","w")
    print('ID'.rjust(4), 'Fitting'.rjust(20), 'Velocity (fpm)'.rjust(15), 'Q (cfm)'.rjust(15), 'DeltaP (in. wg)'.rjust(16),
          'Diameter (in)'.rjust(17), 'Diffuser pressure (in. wg)'.rjust(25))
    print('-'*120)
    orig_stdout = sys.stdout
    file = open("pyductresult.txt", "w+")
    sys.stdout = file
    print('ID'.rjust(4), 'Fitting'.rjust(20), 'Velocity (fpm)'.rjust(15), 'Q (cfm)'.rjust(15), 'DeltaP (in. wg)'.rjust(16),
          'Diameter (in)'.rjust(17), 'Diffuser pressure (in. wg)'.rjust(25))
    print('-'*120)
    sys.stdout = orig_stdout
    for fitting in fittings:
        if fitting['type'] != 'air_handling_unit':
            # TODO: @sziske needs to fix this
            # velocity = 5
            velocity = fitting['flow'] / (np.pi * (fitting['size'] / 12) * (fitting['size'] / 12))
        else:
            velocity = 0.0
            fitting['size'] = 0.0
        if fitting['pdrop'] is None:
            fitting['pdrop'] = 0.0
        if fitting['size'] is None:
            fitting['size'] = 0.0
        if fitting['diffuser_psum'] is None:
            print(repr(fitting['ID']).rjust(4), fitting['type'].rjust(20), ("%.3f" % velocity).rjust(15),
                  ("%.1f" % fitting['flow']).rjust(15), ("%.3f" % fitting['pdrop']).rjust(16),
                  ("%.3f" % fitting['size']).rjust(17))
            orig_stdout = sys.stdout
            sys.stdout = file
            print(repr(fitting['ID']).rjust(4), fitting['type'].rjust(20), ("%.3f" % velocity).rjust(15),
                  ("%.1f" % fitting['flow']).rjust(15), ("%.3f" % fitting['pdrop']).rjust(16),
                  ("%.3f" % fitting['size']).rjust(17))
            sys.stdout = orig_stdout
        else:
            print(repr(fitting['ID']).rjust(4), fitting['type'].rjust(20), ("%.3f" % velocity).rjust(15),
                  ("%.1f" % fitting['flow']).rjust(15), ("%.3f" % fitting['pdrop']).rjust(16),
                  ("%.3f" % fitting['size']).rjust(17), ("%.3f" % fitting['diffuser_psum']).rjust(25))
            orig_stdout = sys.stdout
            sys.stdout = file
            print(repr(fitting['ID']).rjust(4), fitting['type'].rjust(20), ("%.3f" % velocity).rjust(15),
                  ("%.1f" % fitting['flow']).rjust(15), ("%.3f" % fitting['pdrop']).rjust(16),
                ("%.3f" % fitting['size']).rjust(17), ("%.3f" % fitting['diffuser_psum']).rjust(25))
            sys.stdout = orig_stdout

    file.close()

def print_fitting(f):
    print(' ', int(f['ID']), ' ', end='')
    print(f['type'], ' ', end='')
    if f['IDup'] is not None:
        print('connects to: ', f['IDup'], end='')
    if f['BranchUP'] is not None:
        print('-', f['branchUp'])
    else:
        print('\n', end='')
    if f['length'] is not None:
        print('    length: ', f['length'])
    if f['IDdownMain'] is not None:
        print('    IDdownMain: ', int(f['IDdownMain']))
    if f['IDdownBranch'] is not None:
        print('    IDdownBranch', int(f['IDdownBranch']))
    if f['flow'] is not None:
        print('    flow: ', f['flow'])
    if f['fandist'] is not None:
        print('    fandist: ', f['fandist'])
    if f['size'] is not None:
        print('    size: ', f['size'])
    if f['pdrop'] is not None:
        print('    pdrop: ', f['pdrop'])
    if f['pdropMain'] is not None:
        print('    pdropMain: ', f['pdropMain'])
    if f['pdropBranch'] is not None:
        print('    pdropBranch: ', f['pdropBranch'])
    if f['diffuser_psum'] is not None:
        print('    diffuser_psum: ', f['diffuser_psum'])

def print_summary(ducts):
    print('title: ', ducts['title'])
    print('fan_pressure: ', ducts['fan_pressure'])
    print('air_density: ', ducts['air_density'])
    print('roughness: ', ducts['roughness'])
    print('rounding: ', ducts['rounding'])
    fittings = ducts['fittings']
    for f in fittings:
        print_fitting(f)

def calculate(filename):
    file_data = read_input_file(filename)
    ducts = process_keywords(file_data)
    fittings = ducts['fittings']
    make_connections(fittings)
    setup_flowrates(fittings)
    setup_fan_distances(fittings)
    sizing_iterate_nick(ducts)
    print_results(fittings)

if __name__ == '__main__':
    # filename = 'Duct Design Sample Input.txt'
    # calculate(filename)
    print('Please run from main.py')


# -*- coding:utf-8 -*-
'''
主要功能：建立算法库，主要是储存与时间序列曲线匹配的算法程序
@ author: zhyin
@ time: 2019/09/20
'''

import numpy as np
import math
import matplotlib.pyplot as plt

### -*-LC8相关-*- ###

def dn_to_ref(values):
    '''
    完成DN值到反射率值的转化
    :param values: DN值 values
    :return: 反射率值 values
    '''
    values = np.array(values).astype(float)
    #
    multiple_reflectence = 2e-5
    add_reflectance = -0.1
    values[:, 1:] = values[:, 1:] * multiple_reflectence + add_reflectance
    values = values.tolist()
    #
    return values

#
# ANN
class bpNet(object):
    def __init__(self, BpInNode, BpHideNode, BpOutNode, imgNum=1, Greedy=False):
        #
        self.ih_w = np.zeros([BpInNode, BpHideNode])  # 隐含层结点权值
        self.ho_w = np.zeros([BpHideNode, BpOutNode])  # 输出层结点权值
        self.hide_b0 = np.zeros([BpHideNode])  # 隐含层结点阈值
        self.out_b0 = np.zeros([BpOutNode])  # 输出层结点阈值
        self.in_x0 = np.zeros([BpInNode])  # 输入向量值
        self.hide_s0 = np.zeros([BpHideNode])  # 隐含层结点值
        self.out_y0 = np.zeros([BpOutNode])  # 输出层结点值【预测值】
        self.outErr = np.zeros([BpOutNode])
        self.hideErr = np.zeros([BpHideNode])
        #
        self.rate_ih_w = 0.1  # 权值学习率（输入层->隐含层）
        self.rate_ho_w = 0.1  # 权值学习率（隐含层->输出层）
        self.rate_Hide_b0 = 0.1  # 阈值学习率（隐含层）
        self.rate_Out_b0 = 0.1  # 阈值学习率（输出层）
        self.totalErr = 1.0  # 允许的总误差
        #
        self.bpInNode = BpInNode  # 输入层结点数
        self.bpHideNode = BpHideNode  # 隐藏层结点数
        self.bpOutNode = BpOutNode  # 输出层结点数
        #
        if Greedy is False:
            self.out_yd = np.zeros([BpOutNode])  # 真值
        else:
            self.out_yd = np.zeros([imgNum, BpOutNode])

    def genc_randVal(self, low, high):  # 生成[low,high]之间的随机数
        import random
        var = random.random() * (high - low) + low
        return var

    def sigmoidFunc(self, t0):  # 激励函数
        m0 = math.exp(-t0)
        h0 = 1.0 / (1.0 + m0)
        return h0

    def winit(self, weight,dimension):
        if dimension == 1:
            for i in range(np.shape(weight)[0]):
                weight[i] = self.genc_randVal(-0.01, 0.01)
        if dimension == 2:
            for i in range(np.shape(weight)[0]):
                for j in range(np.shape(weight)[1]):
                    weight[i, j] = self.genc_randVal(-0.01, 0.01)

    def bpInitNetFunc(self):
        self.winit(self.ih_w,2)
        self.winit(self.ho_w,2)
        self.winit(self.hide_b0,1)
        self.winit(self.out_b0,1)

    def bpNetTrainFunc(self, InSam, OutSam, imgNum=1, Greedy=False,Ver=1):
        '''
        训练样本
        :param InSam: [8000]一维数组
        :param OutSam: [8000,1]二维数组
        '''
        #
        InSam = np.array(InSam).astype(float)
        OutSam = np.array(OutSam).astype(float)
        trainSample = np.shape(InSam)[0]
        self.totalErr = 0  # 总的误差
        sum = 0  # 和
        z0 = 0  # 激活值
        #
        for isamp in range(trainSample):
            #
            # 输入层 eg:[8000] or [8000,bpInNode]
            if self.bpInNode == 1:
                self.in_x0[0] = InSam[isamp]
            else:
                self.in_x0 = InSam[isamp, :]  # 输入的样本值
            #
            # 输出层
            if Greedy is True and Ver == 1:  # 贪婪算法中输出值 eg[imgNum,isamp] or [imgNum,isamp,bpOutNode]
                if self.bpOutNode == 1:
                    self.out_yd[:,0] = OutSam[:, isamp]
                else:
                    self.out_yd = OutSam[:, isamp, :]  # 期待输出的样本值
            if Greedy is False and Ver == 1:
                if self.bpOutNode == 1:
                    self.out_yd[0] = OutSam[isamp]
                else:
                    self.out_yd = OutSam[isamp, :]  # 期待输出的样本值
            #
            if Greedy is False and Ver ==2:    # Ver2.0版本中真值有所变化:[bpOutNode,isamp]
                self.out_yd = OutSam[:,isamp]  # 期待输出的样本值
            #
            # 正向传播的过程
            # 1）输入层->隐含层1
            for j in range(self.bpHideNode):
                sum = 0.0
                for i in range(self.bpInNode):
                    sum += self.ih_w[i, j] * self.in_x0[i]
                z0 = sum + self.hide_b0[j]
                self.hide_s0[j] = self.sigmoidFunc(z0)  # 隐含层各个单元的输出 1.0/( 1.0 + exp(-z0))
            #
            # 2) 隐含层->输出层
            for j in range(self.bpOutNode):
                sum = 0.0
                for i in range(self.bpHideNode):
                    sum += self.ho_w[i, j] * self.hide_s0[i]
                z0 = sum + self.out_b0[j]
                self.out_y0[j] = self.sigmoidFunc(z0)  # 输出层各个单元的输出 1.0/(1.0 + exp(-z0)
            #
            # 误差反向传播：对于网络中每个输出单元，计算误差值，更新权值
            # 1）输出层->隐含层
            sum = 0.0
            # 计算总均方差
            for j in range(self.bpOutNode):
                #
                # 引入贪婪算法，则此处需要修改
                if Greedy is False:
                    z0 = self.out_yd[j] - self.out_y0[j]
                    self.outErr[j] = z0
                    sum += z0 * z0
                if Greedy is True and Ver == 1:
                    h = 0.0
                    sum = 0.0
                    for m in range(imgNum):
                        h += self.out_yd[m,j] - self.out_y0    # 理论上进行求导
                        sum += (self.out_yd[m,j] - self.out_y0) * (self.out_yd[m,j] - self.out_y0)
                    z0 = h
                    self.outErr[j] = z0
            #
            self.totalErr += sum / 2.0
            #
            for j in range(self.bpOutNode):
                self.outErr[j] = self.outErr[j] * \
                                 self.out_y0[j] * (1.0 - self.out_y0[j])  # 输出层δ2 = ei * θ'(si2)  期望误差*输出层Outy0激励函数导数
                for i in range(self.bpHideNode):
                    self.ho_w[i,j] += self.rate_ho_w * self.outErr[j] * self.hide_s0[i]      # 更新权重
            for j in range(self.bpOutNode):
                self.out_b0[j] += self.rate_Out_b0 * self.outErr[j]     # 更新阈值
            #
            # 2) 隐含层->输入层
            for j in range(self.bpHideNode):
                sum = 0.0
                for i in range(self.bpOutNode):
                    sum += self.outErr[i] * self.ho_w[j,i]
                self.hideErr[j] = sum * self.hide_s0[j] * (1. - self.hide_s0[j]) #隐含层δ1 = （∑out_Er *ho_w） * θ'(si2)  隐含层误差*隐含层Hides0激励函数导数
                for i in range(self.bpInNode):
                    self.ih_w[i,j] += self.rate_ih_w * self.hideErr[j] * self.in_x0[i]  # 更新权重
            for j in range(self.bpHideNode):
                self.hide_b0[j] += self.rate_Hide_b0 * self.hideErr[j]  # 更新阈值

    def bpNetRecognizeFunc(self,testSam):
        '''
        BP神经网络测试样本
        :param testSam: 测试样本
        :return:
        '''
        testSam = np.array(testSam).astype(float)
        # self.totalErr = 0  # 单个样本的误差
        #
        # 输入层 eg:[8000] or [8000,bpInNode]
        if self.bpInNode == 1:
            self.in_x0[0] = testSam
        else:
            self.in_x0 = testSam  # 输入的样本值
        #
        # 输入->隐含层
        for j in range(self.bpHideNode):
            sum = 0.0
            for i in range(self.bpInNode):
                sum += self.ih_w[i, j] * self.in_x0[i]
            z0 = sum + self.hide_b0[j]
            self.hide_s0[j] = self.sigmoidFunc(z0)  # 隐含层各个单元的输出 1.0/( 1.0 + exp(-z0))
        #
        # 2) 隐含层->输出层
        for j in range(self.bpOutNode):
            sum = 0.0
            for i in range(self.bpHideNode):
                sum += self.ho_w[i, j] * self.hide_s0[i]
            z0 = sum + self.out_b0[j]
            self.out_y0[j] = self.sigmoidFunc(z0)  # 输出层各个单元的输出 1.0/(1.0 + exp(-z0)
        #
    pass

#
# demo测试
if __name__ == "__main__":
    file = open('./txt/bpTest.txt','w')
    err = []
    err_time = []
    m_Insample = [[1.78, 1.14], [1.96, 1.18], [1.86, 1.20], [1.72, 1.24], [2.00, 1.26], [2.00, 1.28], [1.96, 1.30],
                  [1.74, 1.36], [1.64, 1.38], [1.82, 1.38], [1.90, 1.38], [1.70, 1.40], [1.82, 1.48], [1.82, 1.54],
                  [2.08, 1.56]]
    m_Outsample = [1, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
    m_Tsample = [[1.80,1.24],[1.84,1.28],[2.04,1.40]]
    # m_Insample = [[1.16,0.116],[1.35,0.104],[1.72,0.078],[1.86,0.107],[1.97,0.136],[2.15,0.082],[2.23,0.125],
    #               [2.48,0.076],[2.79,0.122],[2.85,0.092],[3.07,0.081],[3.45,0.068],[3.59,0.077],[3.80,0.108],
    #               [3.93,0.128],[4.14,0.063],[4.46,0.135],[4.55,0.070],[4.84,0.126],[5.03,0.087]]
    # m_Outsample = [0.502,0.595,0.588,0.662,0.655,0.645,0.736,0.764,0.785,0.792,0.814,0.903,0.931,0.982,0.973,0.981,0.973,0.988,0.969,0.986]
    # m_Tsample = [[1.42,0.086],[2.51,0.071],[3.21,0.107],[4.29,0.096],[5.24,0.065]]
    #
    bp = bpNet(BpInNode=2,BpHideNode=10,BpOutNode=1)
    bp.bpInitNetFunc()

    # 训练
    times = 0
    while bp.totalErr > 0.0001 and times < 10000:
        times += 1
        bp.bpNetTrainFunc(m_Insample,m_Outsample)
        if (times + 1) % 100 == 0:
            file.write('BP %5d DT:%10.5f\n' % ((times+1),bp.totalErr))
            print('BP %5d DT:%10.5f\n' % ((times+1),bp.totalErr))
        err.append(bp.totalErr)
        err_time.append(times+1)
    plt.plot(err_time,err)
    plt.show()
    # 测试
    testSample = np.shape(np.array(m_Tsample))[0]
    for j in range(testSample):
        bp.bpNetRecognizeFunc(np.array(m_Tsample)[j,:])
        print(bp.out_y0)
    pass