from icolos.core.containers.perturbation_map import PerturbationMap
from icolos.core.workflow_steps.pmx.base import StepPMXBase
from icolos.utils.enums.logging_enums import LoggingConfigEnum
from icolos.utils.enums.program_parameters import PMXEnum
from icolos.utils.execute_external.pmx import PMXExecutor
from pydantic import BaseModel
import os

_PE = PMXEnum()
_LE = LoggingConfigEnum()


class StepPMXgentop(StepPMXBase, BaseModel):
    """Fill hybrid topology with B states."""

    def __init__(self, **data):
        super().__init__(**data)
        self._initialize_backend(executor=PMXExecutor)

    def execute(self):
        perturbation_map: PerturbationMap = (
            self.get_workflow_object().get_perturbation_map()
        )

        # generate the B states based on the topology from pmx mutate

        for edge in perturbation_map.get_edges():
            edge_dir = os.path.join(self.work_dir, edge.get_edge_id())
            # overwrite the topology generated by pdb2gmx
            gentop_args = ["-p", "topol.top", "-o", "topol.top"]
            for branch in self.therm_cycle_branches:
                self._logger.log(
                    f"Generating B-state topology for edge {edge.get_edge_id()}, branch {branch}",
                    _LE.DEBUG,
                )
                self._backend_executor.execute(
                    command=_PE.GENTOP,
                    arguments=gentop_args,
                    check=True,
                    location=os.path.join(edge_dir, branch),
                )

            # now we have pmxtop.top in the edge dir


help_string = """
pmx gentop -h
usage: pmx [-h] [-p topol] [-o outfile] [-ff ff] [--split] [--scale_mass]
           [--scale_dih SCALE_DIH] [--norecursive]

This script fills in the B state to a topology file (itp or top) according to
the hybrid residues present in the file. If you provide a top file with
include statemets, by default the script will run through the included itp
files too; this can turned off using the --norecursive flag. You need to use
this script after having mutated a structure file with pmx mutate, and after
having passed that mutated structure through pdb2gmx.

optional arguments:
  -h, --help            show this help message and exit
  -p topol              Input topology file (itp or top). Default is
                        "topol.top"
  -o outfile            Output topology file. Default is "pmxtop.top"
  -ff ff                Force field to use. If -p is a top file, it is not
                        necessary to specify the forcefield, as it will be
                        determined automatically. If -p is an itp file, then
                        -ff is needed, and if not provided a list of available
                        ff will be shown.
  --split               Write separate topologies for the vdW and charge
                        transformations.
  --scale_mass          Scale the masses of morphing atoms so that dummies
                        have a mass of 1.
  --scale_dih SCALE_DIH
                        Scale the dihedrals that have a dummy.
  --norecursive         Whether to fill the B states also for all itp files
                        included in the provided topology file. Default is
                        True. This flag sets it to False.
"""
