# Source code for propertyestimator.properties.binding

    """
    A collection of density physical property definitions.
    """

    from propertyestimator.properties import PhysicalProperty
    from propertyestimator.properties.plugins import register_estimable_property
    from propertyestimator.protocols import coordinates, forcefield, miscellaneous, yank
    from propertyestimator.substances import Substance
    from propertyestimator.workflow.schemas import WorkflowSchema
    from propertyestimator.workflow.utils import ProtocolPath

    @register_estimable_property()
    class HostGuestBindingAffinity(PhysicalProperty):
        """A class representation of a host-guest binding affinity property"""

        @property
        def multi_component_property(self):
            """Returns whether this property is dependant on properties of the
            full mixed substance, or whether it is also dependant on the properties
            of the individual components also.
            """
            return False

        @staticmethod
        def get_default_workflow_schema(calculation_layer, options=None):

            if calculation_layer == "SimulationLayer":
                return HostGuestBindingAffinity.get_default_simulation_workflow_schema(
                    options
                )

            return None

        @staticmethod
        def get_default_simulation_workflow_schema(options=None):
            """Returns the default workflow to use when estimating this property
            from direct simulations.

            Parameters
            ----------
            options: WorkflowOptions
                The default options to use when setting up the estimation workflow.

            Returns
            -------
            WorkflowSchema
                The schema to follow when estimating this property.
            """

            schema = WorkflowSchema(property_type=HostGuestBindingAffinity.__name__)
            schema.id = "{}{}".format(HostGuestBindingAffinity.__name__, "Schema")

            # Initial coordinate and topology setup.
            filter_ligand = miscellaneous.FilterSubstanceByRole("filter_ligand")
            filter_ligand.input_substance = ProtocolPath("substance", "global")

            filter_ligand.component_role = Substance.ComponentRole.Ligand
            # We only support substances with a single guest ligand.
            filter_ligand.expected_components = 1

            schema.protocols[filter_ligand.id] = filter_ligand.schema

            # Construct the protocols which will (for now) take as input a set of host coordinates,
            # and generate a set of charges for them.
            filter_receptor = miscellaneous.FilterSubstanceByRole("filter_receptor")
            filter_receptor.input_substance = ProtocolPath("substance", "global")

            filter_receptor.component_role = Substance.ComponentRole.Receptor
            # We only support substances with a single host receptor.
            filter_receptor.expected_components = 1

            schema.protocols[filter_receptor.id] = filter_receptor.schema

            # Perform docking to position the guest within the host.
            perform_docking = coordinates.BuildDockedCoordinates("perform_docking")

            perform_docking.ligand_substance = ProtocolPath(
                "filtered_substance", filter_ligand.id
            )
            perform_docking.receptor_coordinate_file = ProtocolPath(
                "receptor_mol2", "global"
            )

            schema.protocols[perform_docking.id] = perform_docking.schema

            # Solvate the docked structure using packmol
            filter_solvent = miscellaneous.FilterSubstanceByRole("filter_solvent")
            filter_solvent.input_substance = ProtocolPath("substance", "global")
            filter_solvent.component_role = Substance.ComponentRole.Solvent

            schema.protocols[filter_solvent.id] = filter_solvent.schema

            solvate_complex = coordinates.SolvateExistingStructure("solvate_complex")
            solvate_complex.max_molecules = 1000

            solvate_complex.substance = ProtocolPath(
                "filtered_substance", filter_solvent.id
            )
            solvate_complex.solute_coordinate_file = ProtocolPath(
                "docked_complex_coordinate_path", perform_docking.id
            )

            schema.protocols[solvate_complex.id] = solvate_complex.schema

            # Assign force field parameters to the solvated complex system.
            build_solvated_complex_system = forcefield.BuildSmirnoffSystem(
                "build_solvated_complex_system"
            )

            build_solvated_complex_system.force_field_path = ProtocolPath(
                "force_field_path", "global"
            )

            build_solvated_complex_system.coordinate_file_path = ProtocolPath(
                "coordinate_file_path", solvate_complex.id
            )
            build_solvated_complex_system.substance = ProtocolPath("substance", "global")

            build_solvated_complex_system.charged_molecule_paths = [
                ProtocolPath("receptor_mol2", "global")
            ]

            schema.protocols[
                build_solvated_complex_system.id
            ] = build_solvated_complex_system.schema

            # Solvate the ligand using packmol
            solvate_ligand = coordinates.SolvateExistingStructure("solvate_ligand")
            solvate_ligand.max_molecules = 1000

            solvate_ligand.substance = ProtocolPath("filtered_substance", filter_solvent.id)
            solvate_ligand.solute_coordinate_file = ProtocolPath(
                "docked_ligand_coordinate_path", perform_docking.id
            )

            schema.protocols[solvate_ligand.id] = solvate_ligand.schema

            # Assign force field parameters to the solvated ligand system.
            build_solvated_ligand_system = forcefield.BuildSmirnoffSystem(
                "build_solvated_ligand_system"
            )

            build_solvated_ligand_system.force_field_path = ProtocolPath(
                "force_field_path", "global"
            )

            build_solvated_ligand_system.coordinate_file_path = ProtocolPath(
                "coordinate_file_path", solvate_ligand.id
            )
            build_solvated_ligand_system.substance = ProtocolPath("substance", "global")

            schema.protocols[
                build_solvated_ligand_system.id
            ] = build_solvated_ligand_system.schema

            # Employ YANK to estimate the binding free energy.
            yank_protocol = yank.LigandReceptorYankProtocol("yank_protocol")

            yank_protocol.thermodynamic_state = ProtocolPath(
                "thermodynamic_state", "global"
            )

            yank_protocol.number_of_iterations = 2000
            yank_protocol.steps_per_iteration = 500
            yank_protocol.checkpoint_interval = 10

            yank_protocol.verbose = True

            yank_protocol.force_field_path = ProtocolPath("force_field_path", "global")

            yank_protocol.ligand_residue_name = ProtocolPath(
                "ligand_residue_name", perform_docking.id
            )
            yank_protocol.receptor_residue_name = ProtocolPath(
                "receptor_residue_name", perform_docking.id
            )

            yank_protocol.solvated_ligand_coordinates = ProtocolPath(
                "coordinate_file_path", solvate_ligand.id
            )
            yank_protocol.solvated_ligand_system = ProtocolPath(
                "system_path", build_solvated_ligand_system.id
            )

            yank_protocol.solvated_complex_coordinates = ProtocolPath(
                "coordinate_file_path", solvate_complex.id
            )
            yank_protocol.solvated_complex_system = ProtocolPath(
                "system_path", build_solvated_complex_system.id
            )

            schema.protocols[yank_protocol.id] = yank_protocol.schema

            # Define where the final values come from.
            schema.final_value_source = ProtocolPath(
                "estimated_free_energy", yank_protocol.id
            )

            # output_to_store = WorkflowOutputToStore()
            #
            # output_to_store.trajectory_file_path = ProtocolPath('output_trajectory_path',
            #                                                     extract_uncorrelated_trajectory.id)
            # output_to_store.coordinate_file_path = ProtocolPath('output_coordinate_file',
            #                                                     converge_uncertainty.id, npt_production.id)
            #
            # output_to_store.statistics_file_path = ProtocolPath('output_statistics_path',
            #                                                     extract_uncorrelated_statistics.id)
            #
            # output_to_store.statistical_inefficiency = ProtocolPath('statistical_inefficiency', converge_uncertainty.id,
            #                                                                                     extract_density.id)
            #
            # schema.outputs_to_store = {'full_system': output_to_store}

            return schema