Module bioiain.visualisation.pymol

def quick_display(entity) -> str:
    """
    Displays entity or list of entities with PyMol. Exports entities to ./.temp and saves generated script in the same
    directory as quick_display.pml . Entities are named as N_[entity_id] following input order.
    :param entity: Entity or list of entities.
    :return: Path to the generated script.
    """
    script = PymolScript("quick_display", folder=os.path.join(TEMP_FOLDER,"pml"))

    if type(entity) is not list:
        entity = [entity]
    for n, entity in enumerate(entity):
        name = "{}_{}".format(n, entity.id)
        script.load_entity(entity, name, overwrite=False)
    script.write_script(".temp")
    script.execute()
    return script.path

class PymolScript(object):
    """
    Class to build PyMol scripts from predetermined functions or custom ones.
    :param name: Name of the script. Will de set as a filename. Default is ".temp_pymol_script".
    :return: PymolScript Object.
    """
    def __init__(self, name="temp_pymol_script", folder:str="./pml", tmp_folder=None, pymol_path = "pymol", use_temp=False):
        self.pymol_path = pymol_path
        self._bioiain = "bioiain"
        self.name = name
        self.folder = folder
        if tmp_folder is None:
            tmp_folder = os.path.join(TEMP_FOLDER, self.folder )
        self.tmp_folder=tmp_folder
        os.makedirs(self.folder, exist_ok=True)
        os.makedirs(self.tmp_folder, exist_ok=True)
        self.use_temp = use_temp
        if use_temp:
            self.subfolder =  os.path.join(tmp_folder, self.folder, self.name)
        else:
            self.subfolder = os.path.join(self.folder, self.name)

        os.makedirs(self.subfolder, exist_ok=True)
        self.input = {}
        self.data = {}
        self.commands = []
        self.path = None
        self.session_path = None

    class Command(object):
        """
        Class for commands stored PyMol script.
        :param fun: Name(string) of function to execute. Use .import() to import such function if necessary.
        :param args: Args to pass the function. "strings" -> variables, "'strings'" -> strings.
        :param to: Name of variable to assign the return of the function.
        :param is_cmd: whether the command is within od pymol.cmd.
        :param kwargs: Same as args but with keywords.
        :return: Command object.
        """
        def __init__(self, fun:str, *args, to:str=None, is_cmd=True, **kwargs):
            self.fun = fun
            self.args = args
            self.kwargs = kwargs
            self.to = to
            self.is_cmd = is_cmd
            self.cmd = None

        def __repr__(self):
            if self.cmd is None:
                self.construct_command()
            return self.cmd


        def construct_command(self) -> str:
            """
            Generates final string to append to script. Uses parameters stored in the instance.
            :return: Generated string.
            """
            if type(self.args) == str:
                arg_str = self.args
            else:
                arg_str = ", ".join(self.args)
            kwarg_str = ", ".join([f"{k}={v}" for k, v in self.kwargs.items()])
            c = "{}({},{})".format(self.fun, arg_str, kwarg_str)
            if self.is_cmd:
                c = "cmd."+ c
            if self.to is not None:
                c = "{} = {}".format(self.to, c)
            self.cmd = c
            return self.cmd

    def write_script(self, *args, **kwargs) -> str:
        return self.save(*args, **kwargs)

    def save(self, filename:str=None) -> str:
        """
        Writes the stored commands to a file. The file can be executed from the terminal or run as a PyMol script.
        :param filename: (optional) Path to the file to write to. Uses script name and current wd as default.
        :return: Path to the file.
        """
        if filename is None:
            filename = self.name
        filepath = os.path.join(self.subfolder, filename+".script.pml")
        with open(filepath, "w") as f:
            f.write(f"#!{self.pymol_path}\n\n")
            #f.write(f"try:\n\tcd os.path.dirname(sys.argv[1])\nexcept:\n\tpass\n")
            #f.write("import os\n")
            #f.write(f"os.chdir({self.folder})")

            #f.write("import {} as bi\n\n\n".format(self._bioiain))
            for cmd in self.commands:
                f.write(repr(cmd)+"\n")
        self.path = os.path.abspath(filepath)
        log(1, f"PyMol Session saved at: pymol {self.path}")
        try:
            os.chmod(self.path, 0o755)
        except:
            log("warning", "Could not give exec permissions to script")
        return self.path

    def compile(self, **kwargs):
        return self.execute(compile=True, **kwargs)

    def execute(self, quiet=True, pymol_path=None, full_screen=False, compile=False, use_compiled=True, extra_options="-k"):
        """
        Executes the script on the current thread. Not sure if it is blocking or not.
        """
        if self.path is None:
            self.write_script()
        if pymol_path is None:
            pymol_path = self.pymol_path
        cmd = [pymol_path, extra_options]

        if full_screen:
            cmd.extend(["-x", "-e"])

        if quiet:
            cmd.extend(["-qQ"])

        tmp_session_path = None
        if compile:
            self.session_path = self.path.replace(".script.pml", ".session.pse")
            tmp_session_path = self.path.replace(".script.pml", ".compiler.pml")
            shutil.copy(self.path, tmp_session_path)
            with open(tmp_session_path, "a") as f:
                f.write(f"\ncmd.save('{self.session_path}')")
            cmd.extend(["-c"])
            cmd.extend(["-l", tmp_session_path])
        else:
            if use_compiled and self.session_path is not None:
                cmd.extend([self.session_path])
            else:
                cmd.extend(["-l", self.path])

        logging.log("debug", "$ " + " ".join(cmd))
        try:
            subprocess.run(cmd, cwd=self.subfolder)
            if compile:
                log(1, f"PyMol Session compiled at: pymol {self.session_path}")
                os.remove(tmp_session_path)
        except KeyboardInterrupt:
            logging.log("debug", "\nClosing Pymol...")
        except Exception as e:
            log("error","(PYMOL)", e)


    def add(self, fun, *args, **kwargs) -> Command:
        """
        Generates a command object and adds the command to the script. Can be used to insert custom functions. Beware when adding parameters as strings.
        must be double-quoted ("'string'").
        :param fun: Name(string) of function to execute. Use .import() to import such function if necessary.
        :param args: Args to pass the function. "strings" -> variables, "'strings'" -> strings.
        :param kwargs: Same as args but with keywords.
        :return: Generated Command object.
        """
        c = self.Command(fun, *args, **kwargs)
        self.commands.append(c)
        return c





    @staticmethod
    def _process_sele(sele:str, force_str:bool=False) -> str:
        """
        Adds extra quotes to PyMol selections. Selections are strings within brackets e.g (all).
        :param sele: Selection string.
        :param force_str: Whether to always double-quote input string.
        :return: Double-quoted selection or string if not a selection.
        """
        if sele.startswith("(") and sele.endswith(")") or force_str:
            sele = f"'{sele}'"
        else:
            sele = f"'{sele}'"
        return sele

    @staticmethod
    def _to_str(string):
        return f"'{string}'"


    def print(self, *args, literal=True, **kwargs) -> Command:
        """
        Adds command to print with builtin print.
        :param args: Args to pass to print.
        :param kwargs: Kwargs to pass to print.
        :return: Generated Command object. -> Nothing
        """
        fun = "print"
        if literal:
            new_args = []
            for arg in args:
                new_args.append(self._to_str(repr(arg).replace("'", "\"")))
            args = new_args
        return self.add(fun, *args, is_cmd=False, **kwargs)


    def load(self, path:str, name:str=None, create=False, **kwargs) -> Command:
        """
        Adds command to load file from path.
        :param path: Path to file.
        :param name: Name of created PyMol object.
        :param kwargs:
        :return: Generated Command object -> Unknown.
        """
        if create:
            fun = "create"
        else:
            fun = "load"

        path = relative_path(path, self.subfolder)
        if name is None:
            name = os.path.basename(path).split(".")[0]

        args = f"'{path}'", f"'{name}'"
        return self.add(fun, *args, **kwargs)

    def delete(self, sele="(all)", **kwargs) -> Command:

        fun = "delete"
        sele = self._to_str(sele)
        return self.add(fun, sele, **kwargs)


    def orient(self, sele="(all)", **kwargs) -> Command:

        fun = "orient"
        sele = self._to_str(sele)
        return self.add(fun, sele, **kwargs)

    def center(self, sele="(all)", **kwargs) -> Command:

        fun = "center"
        sele = self._to_str(sele)
        return self.add(fun, sele, **kwargs)

    def show(self, sele="(all)", representation="cartoon", **kwargs):
        fun = "show"
        sele = self._to_str(sele)
        return self.add(fun, self._to_str(representation), sele, **kwargs)

    def hide(self, sele="(all)", representation="everything", **kwargs):
        fun = "hide"
        sele = self._to_str(sele)
        return self.add(fun, self._to_str(representation), sele, **kwargs)

    def load_entity(self, entity, name:str|None=None, overwrite:bool=True) -> Command:
        """
        Adds command to load file from entity. Entity is exported to t/mp/bioiain/pymol as of the cwd.
        :param entity:
        :param name: (optional) Name of created PyMol object. Defaults to the entity id
        :return: Generated Command object -> Unknown.
        """
        if name is None:
            name = entity.data["info"]["name"]
        if not overwrite:
            n = 1
            while name+".pdb" in os.listdir(self.subfolder):
                name = "{}_{}".format(entity.data["info"]["name"], n)
                n += 1

        folder = self.subfolder
        path = entity.export(folder, name, data=True)[0]
        return self.load(path, name)


    def disable(self, sele:str, **kwargs) -> Command:
        """
        Adds Command to disable selection.
        :param sele: Selection string.
        :param kwargs:
        :return: Generated Command object -> Unknown.
        """
        sele = self._to_str(sele)
        fun = "disable"
        return self.add(fun, sele, **kwargs)


    def symmetries(self, obj:str="original", prefix:str="sym", distance:int=6, **kwargs) -> Command:
        fun = "symexp"
        obj = self._to_str(obj)
        args = [self._to_str(prefix), obj, obj, str(distance)]
        return self.add(fun, *args, **kwargs)


    def cell(self, **kwargs) -> Command:
        fun = "show"
        args = "'cell'"
        return self.add(fun, args, **kwargs)

    def group(self, prefix:str="sym", name:str|None=None, also_suffix=False, **kwargs) -> Command:
        fun = "group"
        if also_suffix:
            prefix = "*"+prefix if not prefix.startswith("*") else prefix
        sele = self._to_str(prefix+"*" if  not prefix.endswith("*") else prefix)
        if name is None:
            name = prefix
        args = [self._to_str(name), sele]
        return self.add(fun, *args, **kwargs)

    def align(self, moving:str, fixed:str, fun="align", **kwargs) -> Command:
        sele_fixed = self._to_str(fixed)
        sele_moving = self._to_str(moving)

        args = [sele_moving, sele_fixed]
        return self.add(fun, *args, **kwargs)

    def merge(self, target:str, sele:str, state=-1, **kwargs) -> Command:
        fun = "create"
        sele_target= self._to_str(target)
        sele = self._to_str(sele)

        args = [sele_target, sele]
        return self.add(fun, *args, state=state, **kwargs)


    def color(self, sele:str, color:str|int="black", **kwargs) -> Command:
        fun = "color"
        sele = self._to_str(sele)
        if type(color) is int:
            color = pymol_colours[len(pymol_colours) % (color+1)]
        color = self._to_str(color)
        args = [color, sele]

        return self.add(fun, *args, **kwargs)

    def spectrum(self, sele:str="(all)", spectrum: str = "b", color="rainbow", **kwargs) -> Command:
        fun = "spectrum"
        sele = self._to_str(sele)
        spectrum = self._to_str(spectrum)
        color = self._to_str(color)
        args = [spectrum, color, sele]
        #kwargs["spectrum"] = spectrum

        return self.add(fun, *args, **kwargs)

    def pseudoatom(self, name="tmp", coord=(0,0,0), **kwargs) -> Command:
        fun = "pseudoatom"
        name = self._to_str(name)
        return self.add(fun, name, pos=coord, **kwargs)

    def line(self, name="line", sele1=None, sele2=None, coord1=(0,0,0), coord2=(0,0,0), show_distance=False, **kwargs):
        fun = "distance"
        if isinstance(coord1, PseudoAtom):
            coord1 = coord1.coord
        if isinstance(coord2, PseudoAtom):
            coord2 = coord2.coord
        if sele1 is None:
            sele1 = "tmp1"
            self.pseudoatom(sele1, coord=[float(c) for c in coord1], **kwargs)
        if sele2 is None:
            sele2 = "tmp2"
            self.pseudoatom(sele2, coord=[float(c) for c in coord2], **kwargs)
        args = [self._to_str(name), self._to_str(sele1), self._to_str(sele2)]

        r = self.add(fun, *args, **kwargs)
        if sele1 == "tmp1":
            self.delete(sele1)
        if sele2 == "tmp2":
            self.delete(sele2)
        if not show_distance:
            self.hide(name, "label")
        return r

    def set(self, param, value):
        fun = "set"
        param = self._to_str(param)
        value = self._to_str(value)
        return self.add(fun, param, value)