import phaser
import iotbx.mtz
import scitbx.stl

class TemplateException(phaser.PhaserException):
    pass


class Match(object):
    """
    Matches solutions and symmetry equivalents
    """

    AVERAGE_RESOLUTION = 2.5
    MINIMAL_MODEL_ERROR = 0.6

    def __init__(self, pdb_file_for, unit_cell):

        input = phaser.ENSE()

        for ( name, pdb_file ) in pdb_file_for.items():
            input.addENSE_PDB_RMS( name, pdb_file, self.MINIMAL_MODEL_ERROR )

        self._extent_for = scitbx.stl.map.stl_string_double()
        self._centre_for = phaser.map_string_dvect3()
        self._symm_eulers_for = phaser.map_string_vect_of_dvect3()

        defaults = phaser.data_ptgrp()

        for ( modlid,data_pdb ) in input.PDB.items():
            data_pdb.setup_molecules( defaults )
            self._extent_for[ modlid ] = data_pdb.mean_rad_prpt()
            self._centre_for[ modlid ] = data_pdb.centre()
            self._symm_eulers_for[ modlid ] = data_pdb.symm_eulers()

        data_a = phaser.DataA(
            unit_cell,
            phaser.data_refl( [ ( 0, 0, 0 ) ], [ 0 ], [ 0 ] ),
            phaser.data_resharp()
            )
        self._result = phaser.ResultMR()
        self._result.setDataA( data_a )

        self._known_ensembles = pdb_file_for.keys()


    def is_template_valid(self, template):

        if len( template.KNOWN ) == 1:
            return template.KNOWN[0].getModlid() in self._known_ensembles

        return False


    def check_templates(self, templates):

        # Check templates
        if not all( [ self.is_template_valid( template = t ) for t in templates ] ):
            raise TemplateException, "Invalid template found"

        self._result.setDotSol( templates )
        equiv_count = self._result.calcDuplicateSolutions(
            self._extent_for,
            self._centre_for,
            self._symm_eulers_for,
            self.AVERAGE_RESOLUTION,
            []
            )

        if 0 < equiv_count:
            raise TemplateException, "Non-unique template set"


    def disassemble(self, mr_set):

        singles = []

        for mr_ndim in mr_set.KNOWN:
            if mr_ndim.getModlid() not in self._known_ensembles:
                raise TemplateException, "Unknown ensemble: %s" % mr_ndim.getModlid()

            single_set = phaser.mr_set(
                mr_set.ANNOTATION,
                [ mr_ndim ],
                [],
                mr_set.TF,
                mr_set.TFZ,
                mr_set.LLG,
                mr_set.ORIG_LLG,
                mr_set.R,
                mr_set.ORIG_R,
                mr_set.PAK,
                mr_set.ORIG_NUM,
                mr_set.KEEP,
                mr_set.EQUIV,
                [],
                phaser.MapCoefs(),
                phaser.map_str_float1D(),
                "P 1"
                )
            singles.append( single_set )

        return singles


    def match_single_copies(self, mr_set, templates):

        # Check consistency
        self.check_templates( templates = templates )
        single_sets = self.disassemble( mr_set = mr_set )

        # Save original numbering for templates
        original_template_nums = [ s.ORIG_NUM for s in templates ]
        matched_pairs = []

        # Renumber templates
        for ( index, t ) in enumerate( templates ):
            t.ORIG_NUM = index

        last_index = len( templates )

        for mr_set in single_sets:
            mr_set.ORIG_NUM = last_index
            solution_set = templates + [ mr_set ]

            self._result.setDotSol( solution_set )
            equiv_count = self._result.calcDuplicateSolutions(
                self._extent_for,
                self._centre_for,
                self._symm_eulers_for,
                self.AVERAGE_RESOLUTION,
                []
                )

            if 0 < equiv_count:
                assert equiv_count == 1
                matched_solutions = self._result.getDotSol()
                equivs = [ sol for sol in matched_solutions if not sol.KEEP ]
                assert len( equivs ) == 1
                assert equivs[0].EQUIV < len( matched_solutions )
                orig_index = matched_solutions[ equivs[0].EQUIV ].ORIG_NUM

                if equivs[0].ORIG_NUM == last_index:
                    # Solution was flagged as equivalent
                    matched_pairs.append(
                        ( templates[ orig_index ], mr_set )
                        )
                else:
                    # A template was flagged as equivalent
                    assert orig_index == mr_set.ORIG_NUM
                    matched_pairs.append(
                        ( templates[ equivs[0].ORIG_NUM ], mr_set )
                        )

        # Restore template numbers
        for ( index, t ) in zip( original_template_nums, templates ):
            t.ORIG_NUM = index

        return matched_pairs


    def assemble(self, matching_pairs):

        if any( [ len( p ) != 2 for p in matching_pairs ] ):
            raise TemplateException, "Invalid matching pairs"

        mr_sets = []

        for index, sets in enumerate( zip( *matching_pairs ) ):
            mr_set = phaser.mr_set(
                "",
                sum( [ s.KNOWN for s in sets ], () ),
                [],
                0,
                0,
                0,
                0,
                0.5,
                0.5,
                0,
                index,
                0,
                0,
                [],
                phaser.MapCoefs(),
                phaser.map_str_float1D(),
                "P 1"
                )
            mr_sets.append( mr_set )

        assert len( mr_sets ) == 2
        return mr_sets


    def match_origins(self, matching_pairs):

        self._result.setDotSol( self.assemble( matching_pairs = matching_pairs ) )
        equiv_count = self._result.calcDuplicateSolutions(
            self._extent_for,
            self._centre_for,
            self._symm_eulers_for,
            self.AVERAGE_RESOLUTION,
            []
            )

        return equiv_count == 1

    @classmethod
    def from_mtz(cls, pdb_file_for, mtz_file):

        mtz = iotbx.mtz.object( mtz_file )

        if len( mtz.crystals() ) == 0:
            raise ValueError, "Bad mtz file"

        return cls(
            pdb_file_for = pdb_file_for,
            unit_cell = mtz.crystals()[0].unit_cell_parameters(),
            )


class EquivalentMatcher(object):
    """
    Matches equivalent mr_sets
    """

    def __init__(self, space_group_hall, unit_cell, resolution = 2.5):

        data_a = phaser.DataA(
            unit_cell,
            phaser.data_refl( [ ( 0, 0, 0 ) ], [ 0 ], [ 0 ] ),
            phaser.data_resharp()
            )
        self._result = phaser.ResultMR()
        self._result.setDataA( data_a )
        self._resolution = resolution
	self._space_group_hall = space_group_hall


    def equivalents_of(self, trial, others, definitions):

        trial_enses = sorted( [ n.getModlid() for n in trial.KNOWN ] )

        if not definitions.is_defined( names = set( trial_enses ) ):
            #raise KeyError, "Ensemble definition missing"
            raise KeyError, "Ensemble definition missing: %s" % ", ".join(
                    e for e in trial_enses if not definitions.is_defined( names = [ e ] )
                )

        equivalents = []

        for partial in others:
            enses = sorted( [ n.getModlid() for n in partial.KNOWN ] )

            if enses != trial_enses:
                continue

            if self.is_equivalent(
                left = partial,
                right = trial,
                definitions = definitions
                ):
                equivalents.append( partial )

        return equivalents


    def is_equivalent(self, left, right, definitions):

        self._result.setDotSol( [ left, right ] )
        equiv_count = self._result.calcDuplicateSolutions(
            definitions.extent_for(),
            definitions.centre_for(),
            definitions.symm_eulers_for(),
            self._resolution,
            []
            )

        assert equiv_count in [ 0, 1 ]
        return equiv_count == 1


    def overlap(self, left, right, definitions, minimum = 1):

        left_comps = left.KNOWN # Identification by identity
        right_comps = right.KNOWN # Identification by identity
        equivalents = []

        for lc in left_comps:
            equivs = []
            msl = self.as_mr_set( ndims = [ lc ] )

            for rc in right_comps:
                if lc.getModlid() != rc.getModlid():
                    continue

                msr = self.as_mr_set( ndims = [ rc ] )

                if self.is_equivalent( left = msl, right = msr, definitions = definitions ):
                    equivs.append( rc )

            if equivs:
                equivalents.append( ( lc, equivs ) )

        solutions = []

        for count in range( len( equivalents ), minimum - 1, -1 ):
            solutions = self.matching_combinations(
                equivalents = equivalents,
                count = count,
                definitions = definitions
                )

            if solutions:
                break

        return [
            self.repackage( match = m, lefts = left_comps, rights = right_comps )
            for m in solutions
            ]


    def contains(self, larger, smaller, definitions):

        size = len( smaller.KNOWN )

        if self.overlap(
            left = larger,
            right = smaller,
            definitions = definitions,
            minimum = size
            ):
            return True

        else:
            return False


    def matching_combinations(self, equivalents, count, definitions):

        import itertools
        composites = []

        for combination in itertools.combinations( equivalents, count ):
            ( left, rights ) = zip( *combination )
            ms_left = self.as_mr_set( ndims = left )

            for structure in itertools.product( *rights ):
                if len( set( structure ) ) != count:
                    continue

                ms_right = self.as_mr_set( ndims = structure )

                if self.is_equivalent( left = ms_left , right = ms_right, definitions = definitions ):
                    composites.append( ( left, structure ) )

        return composites


    @classmethod
    def from_mtz(cls, mtz_file):

        mtz = iotbx.mtz.object( mtz_file )

        if len( mtz.crystals() ) == 0:
            raise ValueError, "Bad mtz file"

        return cls(
            unit_cell = mtz.crystals()[0].unit_cell_parameters(),
            space_group_hall = mtz.space_group_info().type().hall_symbol(),
            )


    @staticmethod
    def repackage(match, lefts, rights):

        l_com = set( match[0] )
        r_com = set( match[1] )
        return (
            match,
            [ s for s in lefts if s not in l_com ],
            [ s for s in rights if s not in r_com ],
            )


    def as_mr_set(self, ndims):

        return phaser.mr_set(
            "",
            ndims,
            [],
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            [],
            phaser.MapCoefs(),
            phaser.map_str_float1D(),
            self._space_group_hall
            )


class EnsembleDefinition(object):
    """
    Store ensemble definitions
    """

    def __init__(self):

        self._extent_for = scitbx.stl.map.stl_string_double()
        self._centre_for = phaser.map_string_dvect3()
        self._symm_eulers_for = phaser.map_string_vect_of_dvect3()
        self._names = set()


    def insert(self, name, input):

        defaults = phaser.data_ptgrp()

        for ( modlid,data_pdb ) in input.PDB.items():
            data_pdb.setup_molecules( defaults )
            self._extent_for[ modlid ] = data_pdb.mean_rad_prpt()
            self._centre_for[ modlid ] = data_pdb.centre()
            self._symm_eulers_for[ modlid ] = data_pdb.symm_eulers()

        self._names.add( name )


    def is_defined(self, names):

        return self._names.issuperset( names )


    def knowns(self):

        return self._names


    def unknowns(self, names):

        return set( names ).difference( self._names )


    def define(self, name, identity_for = {}, rms_for = {}):

        input = phaser.ENSE()

        for ( p, i ) in identity_for.items():
            input.addENSE_PDB_ID( name, p, i )

        for ( p, r ) in rms_for.items():
            input.addENSE_PDB_RMS( name, p, r )

        self.insert( name = name, input = input )


    def extent_for(self):

        return self._extent_for


    def centre_for(self):

        return self._centre_for


    def symm_eulers_for(self):

        return self._symm_eulers_for


    def __getstate__(self):

        odict = self.__dict__.copy()
        del odict[ "_extent_for" ]
        del odict[ "_centre_for" ]
        del odict[ "_symm_eulers_for" ]
        return odict


    def __setstate__(self, odict):

        self.__dict__.update( odict )
        self._extent_for = scitbx.stl.map.stl_string_double()
        self._centre_for = phaser.map_string_dvect3()
        self._symm_eulers_for = phaser.map_string_vect_of_dvect3()