import itertools

import numpy as np, enum, abc, scipy.sparse as sp
from McUtils.Combinatorics import SymmetricGroupGenerator, PermutationRelationGraph
import McUtils.Numputils as nput
from ..BasisReps import BasisStateSpace, BasisMultiStateSpace, SelectionRuleStateSpace, BraKetSpace, HarmonicOscillatorProductBasis

__all__ = [
    "DegeneracySpec",
    "DegenerateMultiStateSpace"
]

__reload_hook__ = ["..BasisReps"]

class DegenerateSpaceInputFormat(enum.Enum):
    Groups = "groups"
    QuantaSpecRules = "nT"
    Polyads = "polyads"
    Pairs = "pairs"
    StrongCouplings = "wfc_threshold"
    EnergyCutoff = "energy_window"
    Callable = "callable"
    MartinTest = "martin_threshold"

class DegeneracySpec(metaclass=abc.ABCMeta):
    """
    Provides a container for specifying degeneracies
    in a way that can be cleanly canonicalized
    """

    application_order = 'pre'
    group_filter = None
    def __init__(self,
                 application_order=None, group_filter=None, energy_cutoff=2.25e-3,
                 test_modes=None,
                 maximize_filtered_groups=True,
                 decoupling_overide=100,
                 extra_groups=None,
                 inconsistent_polyads=None
                 ):
        if application_order is None:
            application_order = self.application_order
        self.application_order = application_order
        self.energy_cutoff = energy_cutoff
        if group_filter is None:
            group_filter = self.group_filter
        self.group_filter = group_filter
        self.test_modes = test_modes
        self.maximize_filtered_groups = maximize_filtered_groups
        self.decoupling_overide = decoupling_overide
        self.extra_groups = extra_groups
        self.inconsistent_polyads = inconsistent_polyads

    repr_opts = ['energy_cutoff']
    def __repr__(self):
        return "{}({})".format(type(self).__name__, ", ".join('{}={}'.format(k, getattr(self, k)) for k in self.repr_opts))
    @classmethod
    def merge_state_blocks(cls, state_blocks):
        return [g[1] for g in PermutationRelationGraph.make_relation_graph(state_blocks)[0]]
    def get_degenerate_group_filter(self, solver, corrs=None, threshold=None):
        group_filter = self.group_filter
        if isinstance(group_filter, str) and group_filter == 'default':
            group_filter = None
        if group_filter is None:
            zero_order_energy_cutoff = self.energy_cutoff
            test_modes = self.test_modes
            if test_modes is None:
                test_modes = solver.high_frequency_modes
            if len(test_modes) == 0:
                test_modes = None
            group_filter = dict(
                corrections=corrs,
                energy_cutoff=zero_order_energy_cutoff,
                energies=solver.zero_order_energies,
                threshold=threshold,
                target_modes=test_modes,
                maximize_groups=self.maximize_filtered_groups,
                decoupling_overide=self.decoupling_overide,
                extra_groups=self.extra_groups
            )
        elif isinstance(group_filter, str) and group_filter == 'unfiltered':
            group_filter = None
        return group_filter

    format = None
    @classmethod
    def get_format_mapping(cls):
        return {
            DegenerateSpaceInputFormat.Groups: GroupsDegeneracySpec,
            DegenerateSpaceInputFormat.QuantaSpecRules: TotalQuantaDegeneracySpec,
            DegenerateSpaceInputFormat.Polyads: PolyadDegeneracySpec,
            DegenerateSpaceInputFormat.StrongCouplings: StronglyCoupledDegeneracySpec,
            DegenerateSpaceInputFormat.EnergyCutoff: EnergyCutoffDegeneracySpec,
            DegenerateSpaceInputFormat.Callable: CallableDegeneracySpec,
            DegenerateSpaceInputFormat.MartinTest: MartinTestDegeneracySpec,
        }

    @classmethod
    def get_default_spec(cls):
        return StronglyCoupledDegeneracySpec()
    @classmethod
    def from_spec(cls, spec, format=None, **kwargs):
        if spec is False:
            spec = None
        if spec is None:
            return None
        elif (
                isinstance(spec, str) and spec == 'auto'
                or spec is True
        ):
            return cls.get_default_spec()

        if format is None:
            (spec, opts), format = cls.infer_format(spec)
            if opts is not None:
                kwargs = dict(opts, **kwargs)
        res_cls = cls.get_format_mapping()[format]

        return res_cls(spec, **kwargs)

    @staticmethod
    def _key(obj, key):

        try:
            val = getattr(obj, key)
        except AttributeError:
            val = None
        if val is None:
            try:
                val = obj[key]
            except (TypeError, KeyError, AttributeError, IndexError):
                val = None
        return val


    @classmethod
    def infer_format(cls, spec):

        if spec is None:
            return (None, None), None

        fmt = cls._key(spec, 'format')
        if fmt is not None:
            return spec, fmt

        for fmt in DegenerateSpaceInputFormat:
            key = fmt.value
        #     ['states', DegenerateSpaceInputFormat.Groups],
        #     ['nT', DegenerateSpaceInputFormat.QuantaSpecRules],
        #     ['energy_cutoff', DegenerateSpaceInputFormat.EnergyCutoff],
        #     ['polyads', DegenerateSpaceInputFormat.Polyads],
        #     ['couplings', DegenerateSpaceInputFormat.StrongCouplings],
        #     ['callable', DegenerateSpaceInputFormat.Callable]
        # ]:
            val = cls._key(spec, key)
            if val is not None:
                opts = cls._key(spec, 'options')
                if opts is None and isinstance(spec, dict):
                    opts = spec.copy()
                    del opts[key]
                return (val, opts), fmt
        else:
            for fmt,kls in cls.get_format_mapping().items():
                new_spec = kls.canonicalize(spec)
                if new_spec is not None:
                    return (new_spec, None), fmt
            else:
                return (None, None), None

    @abc.abstractmethod
    def get_groups(self, input_states, solver=None, **kwargs):
        """
        :param solver:
        :type solver:
        :param input_states:
        :type input_states:
        :return:
        :rtype:
        """

        raise NotImplementedError("abstract interface")

    @staticmethod
    def get_group_polyad_relation(exc):
        rows, cols = np.triu_indices(len(exc), k=1)
        diffs = exc[cols] - exc[rows]
        negs = diffs < 0
        a = diffs
        b = np.zeros_like(a)
        b[negs] = -diffs[negs]
        a[negs] = 0
        return np.moveaxis(np.array([a, b]), 0, 1)

    @classmethod
    def get_polyad_pairs_from_polyad_specs(cls, polyads):
        return [
            [p1, p2]
            for polyad in polyads
            for i,p1 in enumerate(polyad)
            for p2 in polyad[i+1:]
        ]
    def get_polyad_pairs(self, input_states, solver=None, **kwargs):
        """
        :param solver:
        :type solver:
        :param input_states:
        :type input_states:
        :return:
        :rtype:
        """
        groups = self.get_groups(input_states, solver=solver, **kwargs)
        if len(groups) == 0:
            base_polyads = []
        else:
            base_polyads = np.unique(
                np.concatenate(
                    [
                        self.get_group_polyad_relation(g)
                        for g in groups
                    ],
                    axis=0
                ),
                axis=0
            )

        # filter out anything that just adds one quantum
        base_polyads = [
            b for b in base_polyads
            if (np.sum(np.abs(b[0])) > 0 and np.sum(np.abs(b[1])) > 0)
        ]

        return self.get_polyad_pairs_from_polyad_specs(base_polyads)

    @classmethod
    @abc.abstractmethod
    def canonicalize(cls, spec):
        raise NotImplementedError("abstract interface")

    def get_group_filter(self, **kwargs):
        return DegenerateMultiStateSpace.construct_filer(**kwargs)

class EnergyCutoffDegeneracySpec(DegeneracySpec):
    """

    """

    format = DegenerateSpaceInputFormat.EnergyCutoff
    def __init__(self, cutoff, **opts):
        super().__init__(**opts)
        self.cutoff = cutoff

    @classmethod
    def canonicalize(cls, spec):
        return isinstance(spec, (int, float, np.integer, np.floating))

    @classmethod
    def _group_states_by_energy_cutoff(cls, H0, states, cutoff):
        """
        :type H: Iterable[SparseArray]
        :type states: BasisStateSpace
        :type total_state_space: BasisMultiStateSpace
        :type cutoff: float
        :rtype: Iterable[BasisStateSpace]
        """
        # we look for states with energies within a range...
        # so initially we pull the sets of energies

        diag_inds = BraKetSpace(states, states)
        energies = H0[diag_inds, diag_inds]
        degenerate_groups = []
        # then we look through the input states
        for n, e in enumerate(energies):
            # we only want to apply this once per degenerate group
            # NOTE: this is a path to subtlety, since
            #   if state a is within 50 cm^-1 of state b, and state b is within of c,
            #   you might argue a and c are degenerate
            #   we are wagering that states are distinct _enough_ such that this is not
            #   an issue, but if it is a different strategy will be required
            if all(n not in d for d in degenerate_groups):
                e_diffs = np.abs(energies - e)
                inds = np.where(e_diffs < cutoff)[0]
                degenerate_groups.append(set(inds))
        # raise Exception(degenerate_groups)
        degenerate_groups = [states.take_subspace(np.array(list(d), dtype=int)) for d in degenerate_groups]
        return degenerate_groups

class MartinTestDegeneracySpec(DegeneracySpec):
    application_order = 'mid'
    group_filter = 'default'
    def __init__(self, threshold=4.6e-6, test_energy_window=4.6e-3, convert=True, frequencies=None, **opts):
        super().__init__(**opts)
        if convert and threshold > 1e-2: # help coerce it into Hartrees
            threshold = threshold / 219475 # only need a rough conversion...
        self.threshold = threshold
        self.window = test_energy_window
        self.frequencies = frequencies
        self._test_groups = None
        self._states = None
        self._basis = None
        self._matrix = None

    repr_opts = ['energy_cutoff', 'threshold']
    def prep_states(self, states:BasisStateSpace):
        state_list = states.excitations
        zo_engs = np.dot(state_list, self.frequencies)
        extended_states = states.basis.operator('x', 'x', 'x').get_transformed_space(states)
        key_inds = []
        exc_groups = []
        for i, (s, e0, exc) in enumerate(zip(state_list, zo_engs, extended_states)):
            e1 = np.dot(exc.excitations, self.frequencies)
            eng_diffs = e1 - e0
            pos = np.where(np.abs(eng_diffs < self.window))
            if len(pos) > 0 and len(pos[0]) > 0:
                key_inds.append(i)
                exc_groups.append(exc.take_subspace(pos[0]))
        if len(key_inds) > 0:
            self._test_groups = SelectionRuleStateSpace(
                states.take_subspace(key_inds),
                exc_groups
            )
            new = states.union(self._test_groups.to_single().take_unique())
        else:
            new = states
            self._test_groups = []
        return new
    def get_coupled_spaces(self, input_states:BasisStateSpace, solver=None):
        # raise NotImplementedError("...")
        if self._test_groups is None:
            self._test_groups = self.prep_states(input_states)
        elif len(self._test_groups) == 0:
            return []
        if solver is None:
            raise ValueError("need a solver with fully evaluated perturbations first...")
        total_basis = solver.flat_total_space
        # engs = solver.zero_order_energies
        h1 = solver.representations[1]
        groups = self._test_groups #type:SelectionRuleStateSpace
        keys = groups.representative_space
        key_inds = total_basis.find(keys)
        spaces = []
        martin_inds = [[], []]
        martin_vals = []
        for i, (k, subspace) in enumerate(zip(key_inds, groups.flat)):
            e0 = np.dot(keys.excitations[i], self.frequencies)
            e1 = np.dot(subspace.excitations, self.frequencies)
            p = total_basis.find(subspace)
            e_diffs = np.abs(e1 - e0)
            repr_elems = h1[k, p].asarray()
            # diffs = subspace.excitations - keys.excitations[np.newaxis, k]
            # diff_sums = np.sum(np.abs(diffs), axis=1)
            # raise Exception(diffs, keys.excitations[np.newaxis, k], subspace.excitations)
            # weights = 1
            test_val = (repr_elems**4)/(e_diffs**3)
            # if (keys.excitations[i] == np.array([0, 0, 0, 1, 0, 0, 0, 0, 0])).all():
            #     raise Exception(test_val*219465, repr_elems, e_diffs, subspace.excitations)
            # if (keys.excitations[i] == np.array([0, 0, 0, 0, 0, 0, 1, 0, 0])).all():
            #     raise Exception(test_val*219465, repr_elems, e_diffs, subspace.excitations)
            martin_inds[0].extend([i]*len(subspace))
            martin_inds[1].extend(p)
            martin_vals.extend(test_val)
            pos = np.where(test_val > self.threshold)
            if len(pos) > 0 and len(pos[0]) > 0:
                # print(pos[0], subspace.excitations)
                new = np.concatenate([
                        keys.excitations[i:i+1],
                        subspace.excitations[pos[0],]#.take_subspace(pos[0]).excitations
                    ], axis=0)
                # print(keys.excitations[i:i+1], new)
                new_sums = np.sum(new, axis=1)
                new = new[new_sums > 0]
                spaces.append(new)

        self._states = keys
        self._basis = solver.flat_total_space
        self._matrix = nput.SparseArray.from_data(
            (
                np.array(martin_vals),
                tuple(np.array(x) for x in martin_inds)
            ),
            shape=(len(self._states), len(self._basis))
        )
        # raise Exception("...")
        return spaces
    # @classmethod
    def get_groups(self, input_states:BasisStateSpace, solver=None, extra_groups=None, **kwargs):
        groups = self.get_coupled_spaces(input_states, solver=solver)
        indexer = SymmetricGroupGenerator(input_states.ndim)
        if extra_groups is None:
            extra_groups = self.extra_groups
        if extra_groups is None:
            extra_groups = []
        extra_groups = [(indexer.to_indices(g), np.asanyarray(g)) for g in extra_groups]
        groups = PermutationRelationGraph.merge_groups([(indexer.to_indices(g), g) for g in groups] + extra_groups)
        return [g[1] for g in groups]

    @classmethod
    def canonicalize(cls, spec):
        return isinstance(spec.threshold, float)

    def get_group_filter(self, threshold=None, states=None, basis=None, corrections=None, **kwargs):
        return DegenerateMultiStateSpace.construct_filer(
            states=self._states,
            basis=self._basis,
            corrections=self._matrix,
            threshold=self.threshold,
            threshold_step_size=1/219475,
            extra_groups=self.extra_groups,
            **kwargs
        )

class StronglyCoupledDegeneracySpec(DegeneracySpec):
    """

    """
    # application_order = 'post'
    format = DegenerateSpaceInputFormat.StrongCouplings
    default_threshold=.3
    def __init__(self, wfc_threshold=None, state_filter=None, extend_spaces=True, iterations=None,
                 evaluator=None,
                 **opts):
        self.evaluator = evaluator
        self._couplings = {}
        super().__init__(**opts)
        if wfc_threshold is None or isinstance(wfc_threshold, str) and wfc_threshold == 'auto':
            wfc_threshold = self.default_threshold
        self.wfc_threshold = wfc_threshold
        self.state_filter = state_filter
        self.extend_spaces=extend_spaces
        self._iterations = iterations
        self.iterations = iterations

    @property
    def application_order(self):
        return 'post' if self.evaluator is None else 'pre'
    @application_order.setter
    def application_order(self, ord):
        ...

    repr_opts = ['energy_cutoff', 'wfc_threshold']
    def prep_states(self, input_states):
        return input_states

    def identify_strong_couplings(self, solver, corrs):
        degenerate_correction_threshold = self.wfc_threshold
        test_modes = self.test_modes
        if test_modes is None:
            test_modes = solver.high_frequency_modes

            # raise Exception(test_modes, test_freqs * 219465, fundamentals.excitations, where)

        state_filter = self.state_filter
        zero_order_energy_cutoff = self.energy_cutoff
        if state_filter is None:
            state_filter = lambda state, couplings: corrs.default_state_filter(state, couplings,
                                                                               energy_cutoff=zero_order_energy_cutoff,
                                                                               energies=solver.zero_order_energies,
                                                                               basis=solver.flat_total_space,
                                                                               target_modes=test_modes
                                                                               )
        sc = corrs.find_strong_couplings(threshold=degenerate_correction_threshold, state_filter=state_filter)
        return sc

    def get_input_state_couplings(self, input_states):
        inds = input_states.indices
        needs_couplings = [
            n for n,i in enumerate(inds)
            if i not in self._couplings
        ]
        if len(needs_couplings) > 0:
            needs_coupling_states = input_states.take_subspace(needs_couplings)
            wfcs = self.evaluator.get_test_wfn_corrs(needs_coupling_states, self.energy_cutoff)
            if wfcs is not None:
                for input_state, final_state, corrs in zip(wfcs.initial_states.indices, wfcs.final_states, wfcs.corrections):
                    coupling_pos = np.where(np.abs(corrs) > self.wfc_threshold)
                    if len(coupling_pos) > 0 and len(coupling_pos[0]) > 0:
                        coupling_pos = coupling_pos[1]
                        self._couplings[input_state] = final_state.take_subspace(coupling_pos)
            for i in inds:
                if i not in self._couplings:
                    self._couplings[i] = None
        return {
            i:self._couplings[i]
            for i in inds
            if self._couplings[i] is not None
        }

    def get_groups(self, input_states, couplings=None, solver=None, extra_groups=None, **kwargs):
        """
        :param input_states:
        :type input_states:
        :param solver:
        :type solver:
        :return:
        :rtype:
        """
        if self.evaluator is not None:
            couplings = self.get_input_state_couplings(input_states)
            # if len(couplings) == 0:
            #     return None
        else:
            if couplings is None:
                raise ValueError("need couplings")
        if extra_groups is None:
            extra_groups = self.extra_groups
        return self.get_strong_coupling_space(input_states, couplings, extra_groups=extra_groups)

    @classmethod
    def canonicalize(cls, spec):
        if isinstance(spec, (int, float, np.integer, np.floating)):
            return spec
        else:
            try:
                return {k:GroupsDegeneracySpec._validate_grp(v) for k,v in spec.items()}
            except (np.VisibleDeprecationWarning, AttributeError, StopIteration):
                return None

    @classmethod
    def get_strong_coupling_space(cls, states: BasisStateSpace, couplings: dict, extra_groups=None):
        indexer = SymmetricGroupGenerator(states.ndim)
        groups = [None] * len(states.indices)
        for n, i in enumerate(states.indices):
            if i in couplings:
                groups[n] = [i] + list(couplings[i].indices)
            else:
                groups[n] = [i]

        if extra_groups is None:
            extra_groups = []
        extra_groups = [(indexer.to_indices(g), np.asanyarray(g)) for g in extra_groups]

        # print("premerge groups:", groups)

        groups = PermutationRelationGraph.merge_groups(
            [(g, indexer.from_indices(g)) for g in groups] + extra_groups
        )

        # print("????", [g[0] for g in groups]) # need to track down why we're getting a negative index here
                              # presumably from an integer overflow in a int8 or a conversion
                              # of some sort...

        return [g[1] for g in groups]

class GroupsDegeneracySpec(DegeneracySpec):
    """

    """

    format = DegenerateSpaceInputFormat.Groups
    def __init__(self, groups, **opts):
        super().__init__(**opts)
        self.groups = groups

    @staticmethod
    def _validate_grp(grp):
        try:
            b = grp[0]
        except (IndexError, TypeError):
            raise StopIteration("malformatted")
        else:
            return np.asanyarray(grp)
    @classmethod
    def canonicalize(cls, spec):
        try:
            return [cls._validate_grp(s) for s in spec]
        except (np.VisibleDeprecationWarning, AttributeError, StopIteration):
            return None

    def get_groups(self, input_states, solver=None, **kwargs):
        """
        :param solver:
        :type solver:
        :param input_states:
        :type input_states:
        :return:
        :rtype:
        """
        # do I want to canonicalize if group size is too small?
        return self.groups

class PolyadDegeneracySpec(DegeneracySpec):
    """

    """
    format = DegenerateSpaceInputFormat.Polyads
    def __init__(self, polyads,
                 max_quanta=None,
                 iterations=2,
                 require_converged=False,
                 extra_groups=None,
                 extra_polyads=None,
                 full_group_polyads=True,
                 **opts
                 ):
        super().__init__(**opts)
        self.polyads = polyads
        self.max_quanta = max_quanta
        self.iterations = iterations
        self.require_converged = require_converged
        self.extra_groups = extra_groups
        self.full_group_polyads = full_group_polyads

    @staticmethod
    def _validate_rule(pair):
        try:
            if len(pair) != 2:
                raise StopIteration("malformatted")
            b = pair[0]
        except (IndexError, TypeError):
            raise StopIteration("malformatted")
        else:
            if len(pair[0]) == len(pair[1]):
                return pair
            else:
                raise StopIteration("malformatted")
    @classmethod
    def _check_rule(cls, pair):
        try:
            cls._validate_rule(pair)
        except StopIteration:
            return False
        else:
            return True
    @classmethod
    def canonicalize(cls, spec):
        try:
            return [cls._validate_rule(s) for s in spec]
        except (np.VisibleDeprecationWarning, AttributeError, StopIteration):
            return None

    def get_groups(self, input_states, solver=None, **kwargs):
        """
        :param solver:
        :type solver:
        :param input_states:
        :type input_states:
        :return:
        :rtype:
        """
        if isinstance(input_states, BasisStateSpace):
            input_states = input_states.excitations
        return self.get_degenerate_polyad_space(
            input_states,
            self.polyads,
            max_quanta=self.max_quanta,
            iterations=self.iterations,
            require_converged=self.require_converged,
            extra_groups=self.extra_groups
        )
    @classmethod
    def get_degenerate_polyad_space(cls, states, polyadic_pairs,
                                    max_quanta=None, iterations=2,
                                    require_converged=False, extra_groups=None
                                    ):
        """
        Gets degenerate spaces by using pairs of transformation rules to
        take an input state and connect it to other degenerate states

        :param states: the input states
        :type states:
        :param polyadic_pairs: the transformation rules
        :type polyadic_pairs:
        :param max_quanta: the max quanta to allow in connected states
        :type max_quanta:
        :return:
        :rtype:
        """

        if cls._check_rule(polyadic_pairs[0]): # we only have one space
            polyadic_pairs = [polyadic_pairs]
        if isinstance(iterations, int):
            iterations = [iterations]*len(polyadic_pairs)

        groups = None
        for polyad, iterations in zip(polyadic_pairs, iterations):
            graph = PermutationRelationGraph(polyad)
            if groups is None:
                groups = graph.build_state_graph(states,
                                                 extra_groups=extra_groups,
                                                 max_sum=max_quanta,
                                                 max_iterations=iterations,
                                                 raise_iteration_error=require_converged
                                                 )
            else:
                _ = []
                for states in groups:
                    sub = np.concatenate(graph.build_state_graph(states,
                                                     extra_groups=extra_groups,
                                                     max_sum=max_quanta,
                                                     max_iterations=iterations,
                                                     raise_iteration_error=require_converged
                                                     ))
                    _.append(sub)
                groups = _
        groups = [g for g in groups if len(g) > 1]
        return groups



    def get_polyad_pairs(self, input_states, solver=None, full_group_polyads=None, **kwargs):
        if full_group_polyads is None:
            full_group_polyads = self.full_group_polyads
        if full_group_polyads:
            return super().get_polyad_pairs(input_states, solver=solver, **kwargs)
        else:
            return self.get_polyad_pairs_from_polyad_specs(self.polyads)

    @staticmethod
    def _is_polyad_rule(d, n_modes):
        try:
            return (
                    len(d) == 2
                    and len(d[0]) == n_modes
                    and len(d[1]) == n_modes
            )
        except TypeError:
            return False

class TotalQuantaDegeneracySpec(PolyadDegeneracySpec):
    """

    """

    format = DegenerateSpaceInputFormat.QuantaSpecRules
    def __init__(self, n_T_vectors, max_quanta=3, target_modes=None, extra_polyads=None, **opts):
        self.max_quanta = max_quanta
        self.target_modes = target_modes
        self.nt_vecs = [n_T_vectors] if isinstance(n_T_vectors[0], int) else n_T_vectors
        if extra_polyads is None:
            extra_polyads = []
        if len(extra_polyads) > 0 and self._check_rule(extra_polyads[0]):
            extra_polyads = [extra_polyads]
        rules = [
                    self.make_nt_polyad(v, max_quanta=max_quanta) for v in self.nt_vecs
                ] + extra_polyads
        # raise Exception(rules)
        super().__init__(rules, **opts)

    @classmethod
    def reduce_rule(cls, a, b):
        a = a.copy()
        b = b.copy()

        m1 = np.min(a[a>0])
        m2 = np.min(b[b>0])
        if m2 > m1:
            m1, m2 = m2, m1
        if m1 % m2 == 0:
            a = a//m2
            b = b//m2
        reduce_pos = np.where(np.logical_and(a > 0, b > 0))
        if len(reduce_pos) > 0:
            reduce_pos = reduce_pos[0]
            if len(reduce_pos) > 0:
                pairs = np.concatenate([a[reduce_pos][:, np.newaxis], b[reduce_pos][:, np.newaxis]], axis=1)
                min_vals = np.min(pairs, axis=1)
                a[reduce_pos] = a[reduce_pos] - min_vals
                b[reduce_pos] = b[reduce_pos] - min_vals
        return a,b

    @classmethod
    def make_nt_polyad(cls, nt, target_modes=None, max_quanta=3):  # this should be of the form num of quanta give same E
        # we want to choose different subsets of the nT vector to construct rules
        # we'll do this in a kinda dumb way where we force a max number of quanta for these polyad
        # rules and explicitly test states of a given form
        ndim = len(nt)
        nt = np.asanyarray(nt)
        if target_modes is None:
            # target_modes = np.arange(len(nt))
            target_modes = np.where(nt != 0)[0]

        nt = nt[target_modes]
        states = BasisStateSpace.from_quanta(
            HarmonicOscillatorProductBasis(len(nt)),
            list(range(1, max_quanta+1))
        ).excitations

        nts = np.dot(states, nt)
        (keys, groups), _ = nput.group_by(states, nts)
        rules = set()
        for k,g in zip(keys, groups):
            if k > 0:
                for n,i in enumerate(g):
                    for j in g[n+1:]:
                        a, b = cls.reduce_rule(i, j)
                        # insert zeros at dropped positions
                        if len(target_modes) < ndim:
                            _ = np.zeros(ndim, dtype=int)
                            _[target_modes] = a
                            a = _
                            _ = np.zeros(ndim, dtype=int)
                            _[target_modes] = b
                            b = _
                        a = tuple(a)
                        b = tuple(b)
                        rules.add((a, b))
        return list(rules)

class CallableDegeneracySpec(DegeneracySpec):
    """

    """

    format = DegenerateSpaceInputFormat.Callable
    def __init__(self, callable, **opts):
        super().__init__(**opts)
        self.callable = callable

    def get_groups(self, input_states, solver=None, **kwargs):
        """
        :param input_states:
        :type input_states:
        :param solver:
        :type solver:
        :return:
        :rtype:
        """
        return self.callable(input_states, solver=solver, **kwargs)

class DegenerateMultiStateSpace(BasisMultiStateSpace):

    @staticmethod
    def default_group_filter(group,
                             states=None,
                             basis=None,
                             corrections=None,
                             threshold=None,
                             energy_cutoff=None,
                             energies=None,
                             decoupling_overide=10,
                             maximize_groups=True,
                             target_modes=None,
                             threshold_step_size=.1,
                             extra_groups=None
                             ):
        """
        Excludes modes that differ in only one position, prioritizing states with fewer numbers of quanta
        (potentially add restrictions to high frequency modes...?)

        :param input_state:
        :type input_state:
        :param couplings:
        :type couplings:
        :return:
        :rtype:
        """

        if len(group) < 2:
            return group

        if target_modes is not None:
            sub = group.take_subdimensions(target_modes)
            inds = sub.indices
            exc = sub.excitations
        else:
            inds = group.indices
            exc = group.excitations

        if corrections is not None:
            if isinstance(corrections, np.ndarray):
                state_inds = basis.find(group)
                group_inds = state_inds
                found_pos = state_inds < len(states)
            elif hasattr(corrections, 'total_basis'): # PT corrections
                states = corrections.states
                basis = corrections.total_basis
                group_inds = basis.find(group)
                state_inds = states.find(group, dtype=int, missing_val=len(states))
                found_pos = state_inds < len(states)
                kill_spots = (np.arange(len(group_inds))[found_pos], state_inds[found_pos])
                corr_mat = [np.abs(x[:, group_inds].asarray().T) for x in corrections.wfn_corrections[1:]]
                max_corr = None
                for c in corr_mat:
                    c[kill_spots] = 0
                    if max_corr is None:
                        max_corr = np.max(c, axis=1)
                    else:
                        submax = np.max(c, axis=1)
                        p = submax > max_corr
                        max_corr[p] = submax[p]
                corrections = sum(corr_mat)
                if threshold is not None:
                    sorting = np.arange(len(max_corr))
                else:
                    sorting = np.argsort(-max_corr - 1 * (found_pos))  # prioritize states inside requested packet for marginal cases
            elif basis is not None and states is not None:
                group_inds = basis.find(group)
                state_inds = states.find(group, dtype=int, missing_val=len(states))
                found_pos = state_inds < len(states)
                found_states = state_inds[found_pos]
                kill_spots = (np.arange(len(group_inds))[found_pos], found_states)
                corrections = np.abs(corrections[:, group_inds].asarray().T)
                max_corr = None
                corrections[kill_spots] = 0
                if max_corr is None:
                    max_corr = np.max(corrections, axis=1)
                else:
                    submax = np.max(corrections, axis=1)
                    p = submax > max_corr
                    max_corr[p] = submax[p]
                if threshold is not None:
                    sorting = np.arange(len(max_corr))
                else:
                    sorting = np.argsort(-max_corr - 1*found_pos)  # prioritize states inside requested packet for marginal cases

        else:
            sorting = np.argsort(inds)
            corrections = None

        exc = exc[sorting]
        diffs = exc[:, np.newaxis] - exc[np.newaxis, :]  # difference matrix (s, s, m)
        tots = exc[:, np.newaxis] + exc[np.newaxis, :]  # difference matrix (s, s, m)
        diff_sums = np.sum(np.abs(diffs), axis=2)  # (s, s)
        tots_sums = np.sum(np.abs(tots), axis=2)  # (s, s)
        elims = []
        if threshold is not None and corrections is not None:
            if energy_cutoff is None:
                e_vec = None
                bad_pos = np.where(np.logical_and(
                    diff_sums == 1,
                    tots_sums > 0
                ))  # np.logical_or(diff_sums == 1, diff_sums == 0))
            else:
                e_vec = energies[group_inds]
                bad_pos = np.where(np.abs(np.subtract.outer(e_vec, e_vec)) > energy_cutoff)
            if len(bad_pos) > 0 and len(bad_pos[0]) > 0:
                N = len(group_inds)
                base_corr = corrections
                base_vals = base_corr[:, state_inds[found_pos]]
                base_mat = np.zeros((N, N), dtype=float)
                base_mat[:, np.where(found_pos)[0]] = base_vals
                base_mat = (base_mat.T + base_mat) - 2*np.diag(np.diag(base_mat))
                # if decoupling_overide > 0:
                #     truly_bad = base_mat[bad_pos] < decoupling_overide
                # else:
                #     truly_bad = np.full(len(base_mat[bad_pos]), True)
                # # introduce an override so things that are super strongly coupled can remain coupled
                bad_pos = tuple(b for b in bad_pos)
                base_mat[bad_pos] = 0  # refuse to couple problem states
                base_mat[bad_pos[1], bad_pos[0]] = 0  # refuse to couple problem states

                def get_groups(cm):
                    _, labs = sp.csgraph.connected_components(cm, return_labels=True)
                    (grp_keys, groups), _ = nput.group_by(np.arange(N), labs)
                    A = np.zeros((N, N))
                    for g in groups:
                        A[np.ix_(g, g)] = 1
                    return groups, A

                corr_mat = (base_mat > threshold).astype(int)
                groups, A = get_groups(corr_mat)

                target_threshold = threshold
                step_size = threshold_step_size
                while (A[bad_pos] > 0).any() and decoupling_overide > threshold:
                    threshold = threshold + step_size
                    corr_mat = (base_mat > threshold).astype(int)
                    groups, A = get_groups(corr_mat)

                if maximize_groups or (decoupling_overide > 0 and decoupling_overide < threshold):
                    base_thresh = target_threshold
                    if not maximize_groups:
                        base_thresh = decoupling_overide
                    dropped_pos = np.where(np.logical_and(base_mat > base_thresh, base_mat < threshold))
                    if len(dropped_pos) > 0:
                        # figure out if any can be added without causing breakdowns
                        sorting = np.argsort(-base_mat[dropped_pos])
                        dropped_pos = tuple(d[sorting] for d in dropped_pos)
                        for i,j in zip(*dropped_pos): # this could be very expensive...
                            if e_vec is not None:
                                # we can maybe take a shortcut on figuring out if we can extend the group
                                g_i = None
                                g_j = None
                                # first we check if we have a group with `i` or `j` in it
                                for n,g in enumerate(groups):
                                    ix,jx = np.searchsorted(g, [i, j])
                                    if ix < len(g) and g[ix] == i:
                                        g_i = [n, g]
                                    if jx < len(g) and g[jx] == j:
                                        g_j = [n, g]
                                    if g_i is not None and g_j is not None: # need proper graph tools
                                        break
                                else:
                                    # we're just adding to a group, not merging
                                    if g_j is not None:
                                        g_i, g_j = g_j, g_i
                                        i, j = j, i
                                    if g_i is not None:
                                        e_g = e_vec[g_i[1]]
                                        w = max(abs(e_vec[i] - np.min(e_g)), abs(e_vec[i] - np.max(e_g)))
                                        if w < energy_cutoff:
                                            groups[g_i[0]] = np.sort(np.concatenate([g_i[1], [j]]))
                                        continue
                                    # else: # should never matter: would mean states were dropped but are within the energy window

                            corr_mat[i, j] = base_mat[i, j]
                            _, A = get_groups(corr_mat)
                            if (A[bad_pos] > 0).any():
                                corr_mat[i, j] = 0.
                            else:
                                groups = _

                    if (
                            (decoupling_overide > 0 and decoupling_overide < threshold)
                            or extra_groups is not None
                    ):
                        # states with blowups that we need to reinsert
                        if (decoupling_overide > 0 and decoupling_overide < threshold) is not None:
                            need_pos = np.where(np.logical_and(base_mat > decoupling_overide, base_mat < threshold))
                        else:
                            need_pos = ()
                        if extra_groups is not None:
                            for g in extra_groups:
                                try:
                                    new_inds = group.find(g)
                                except IndexError:
                                    pass
                                else:
                                    new_inds = list(itertools.combinations(new_inds, 2))
                                    new_inds = ([x[0] for x in new_inds], [x[1] for x in new_inds])
                                    if len(need_pos) > 0:
                                        need_pos = (
                                            np.concatenate([need_pos[0], new_inds[0]]),
                                            np.concatenate([need_pos[1], new_inds[1]]),
                                        )
                                    else:
                                        need_pos = new_inds
                        if len(need_pos) > 0:
                            sorting = np.argsort(-base_mat[need_pos])
                            need_pos = tuple(d[sorting] for d in need_pos)
                            for i,j in zip(*need_pos):
                                # we try to find a group to insert into
                                g_i = None
                                g_j = None
                                # first we check if we have a group with `i` or `j` in it
                                for n,g in enumerate(groups):
                                    ix,jx = np.searchsorted(g, [i, j])
                                    if ix < len(g) and g[ix] == i:
                                        g_i = [n, g]
                                    if jx < len(g) and g[jx] == j:
                                        g_j = [n, g]
                                    if g_i is not None and g_j is not None:
                                        if g_i[0] == g_j[0]: # nothing needed to do
                                            break
                                        # we move j from its group into i's group
                                        # or vice versa
                                        i_couplings = base_mat[i, g_i[1]]
                                        j_couplings = base_mat[j, g_j[1]]
                                        if np.linalg.norm(j_couplings) > np.linalg.norm(i_couplings):
                                            # swap which gets moved
                                            g_i, g_j = g_j, g_i
                                            i, j = j, i
                                        groups[g_i[0]] = np.sort(np.concatenate([g_i[1], [j]]))
                                        groups[g_j[0]] = np.setdiff1d(g_j[1], [j])
                                        break
                                else:
                                    # we're just adding to a group, not merging
                                    if g_j is not None:
                                        g_i, g_j = g_j, g_i
                                        i, j = j, i
                                    if g_i is not None:
                                        groups[g_i[0]] = np.sort(np.concatenate([g_i[1], [j]]))
                                        continue
                                    else: # I don't think we'll ever get here but...?
                                        groups.append(np.array([i, j]))

                group = [group.take_subspace(g) for g in groups if len(g) > 1]

        else:
            for n in np.arange(len(exc)):
                diff_vec = diff_sums[n]
                bad_pos = np.where(diff_vec == 1)
                if len(bad_pos) > 0:
                    bad_pos = bad_pos[0]
                    if len(bad_pos) > 0:
                        kills = bad_pos[bad_pos > n]
                        elims.extend(sorting[kills])
                        diff_sums[bad_pos, :] = 0

            if len(elims) > 0:
                mask = np.setdiff1d(np.arange(len(exc)), elims)
                group = group.take_subspace(mask)

        return group

    @classmethod
    def construct_filer(cls, spec=None, **kwargs):
        if spec is not None:
            return spec.get_group_filter(**kwargs)
        else:
            return lambda group:cls.default_group_filter(group, **kwargs)

    @classmethod
    def from_spec(cls,
                  degenerate_states,
                  solver=None,
                  full_basis=None,
                  format=None,
                  group_filter=None,
                  log_groups=False,
                  **kwargs
                  ):
        """
        Generates a DegenerateMultiStateSpace object from a number
        of possible specs

        :param solver: the actual applier of the perturbation theory which makes use of the degenerate states
        :type solver: PerturbationTheorySolver
        :return:
        :rtype:
        """

        logger = solver.logger
        states = solver.states #type: BasisStateSpace

        if degenerate_states is not None:
            if not isinstance(degenerate_states, DegeneracySpec):
                degenerate_states = DegeneracySpec.from_spec(degenerate_states, format=format)
            with logger.block(tag="getting degeneracies"):
                deg_states = degenerate_states.get_groups(states, solver=solver, **kwargs)
                if hasattr(group_filter, 'items'):
                    group_filter = cls.construct_filer(**dict(group_filter, spec=degenerate_states))
                new = GroupsDegeneracySpec.canonicalize(deg_states)
                if new is None:
                    raise ValueError("{} returned invalid degenerate subspace spec {}".format(
                        degenerate_states,
                        deg_states
                    ))
                else:
                    degenerate_states = new
                    del deg_states

                degenerate_states = [
                    BasisStateSpace(states.basis, s)#.as_sorted(track_indices=False, track_excitations=False)
                    if not isinstance(s, BasisStateSpace) else s
                    for s in degenerate_states
                ]

                # dds = [d.indices for d in degenerate_states]
                # neg_ps = [i for i,d in enumerate(dds) if np.any(d < 0)]

                # print("DD:", neg_ps, [degenerate_states[i].excitations for i in neg_ps])

                if log_groups:
                    with logger.block(tag="initial degenerate groups:"):
                        for g in degenerate_states:
                            if len(g) > 1:
                                logger.log_print(str(g.excitations).splitlines())
                else:
                    logger.log_print(
                        "{n} degenerate state sets found",
                        n=len([x for x in degenerate_states if len(x) > 1]),
                        # s=[x for x in degenerate_states if len(x) > 1]
                    )

        # build groups of degenerate states for use later
        if degenerate_states is None:
            groups = states.split(1) #[[x] for x in states.indices]  # we're gonna loop through this later so why not destructure now...
        else:
            groups = [None] * len(degenerate_states)
            deg_sets = [(set(d.indices),d) for d in degenerate_states]
            for x in states.indices:
                for i, (d,bs) in enumerate(deg_sets):
                    # if i == 11:
                    #     raise Exception(bs.excitations, bs.indices)
                    if x in d:
                        if groups[i] is None:
                            groups[i] = bs.indices
                        # groups[i].append(x)
                        break
                else:
                    groups.append([x])

        # now turn these into proper BasisStateSpace objects so we can work with them more easily
        ugh = []
        for i,g in enumerate(groups):
            if g is None:
                continue #???
            # g = np.sort(np.array(g))
            if not isinstance(g, BasisStateSpace):
                # if np.any(np.asanyarray(g) < 0):
                #     raise Exception(i, g)
                g = BasisStateSpace(states.basis, np.array(g), mode=BasisStateSpace.StateSpaceSpec.Indices, full_basis=full_basis)
            if group_filter is not None:
                g = group_filter(g)
                if isinstance(g, BasisStateSpace):
                    ugh.append(g)
                else:
                    ugh.extend(g)
            else:
                ugh.append(g)

        for n,x in enumerate(states.indices):
            for g in ugh:
                if x in g.indices:
                    break
            else:
                ugh.append(states.take_subspace([n]))

        # now make a numpy array for initialization
        arrs = np.full(len(ugh), None)
        for i, g in enumerate(ugh):
            arrs[i] = g

        # raise Exception(arrs)

        return cls(arrs)
