#!/usr/bin/env python3
# -----------------------------------------------------------------------------
# Module for controlling fpocket
# Seth Veenbaas
# Weeks Lab, UNC-CH
# 2022
#
# Version 1.3.0
#
# -----------------------------------------------------------------------------

import os
import sys
import subprocess
import shutil
import time
from pymol import cmd


def find_pockets(
    pdb : str,
    chain : str,
    state : int,
    m : float,
    M : float,
    i :int,
    D : float,
    A : int,
    p : float,
    out : str,
    name : str,
    yes : bool,
    ) -> str:
    """Pocket finding pipeline:
        - cleans pdb file to generate rna-only file
        - runs pocket prediction using fpocket
        - manages fpocket output files

    Args:
        pdb (str): Path to input .pdb file.
        chain (str): Chain identifier for desired RNA chain (default='A').
        state (int): Structural state to analyze.
        m (float): Min. a-sphere radius in angstroms (default=3.0).
        M (float): Max. a-sphere radius in angstroms (default=5.7).
        i (int): Min. number of a-spheres per pocket (default=42).
        D (float): a-sphere clustering distance in angstroms (default=1.65).
        out (str): name of fpocket output parent directory name.
        yes (boolean): Overwrite output files and directories with same name.

    Returns:
        str: path to output directory contianing fpocket outputs for analysis
    """

    # Path to clean (ligand-free) pdb file
    pdb_clean = f'{name}_clean.pdb'

    # Makes pdb_clean if it is not already a file
    if not os.path.isfile(pdb_clean):
        clean_pdb(pdb, pdb_clean)

    # Runs fpocket on cleaned pdb file.
    run_fpocket(pdb_clean, name, chain, state, m, M, i, D, A, p)

    # Files fpocket outputs into directories and manages overwriting.
    analysis, yes = file_fpocket(pdb_clean, state, out, yes)
    return analysis.strip('/'), yes

# -----------------------------------------------------------------------------


def clean_pdb(pdb : str, pdb_clean : str) -> None:
    """Cleans a .pdb file input and saves output as a .pdb file.
       Removes not polymer molecules (ligands) and proteins.
       Preserves modified/heteroatom RNA residues.

    Args:
        pdb (str): path to input .pdb file.
        pdb_clean (str): path to output (cleaned) .pdb file.
    """
    cmd.load(pdb)
    cmd.alter('polymer', 'type="ATOM"')
    cmd.save(pdb, state='0')
    time.sleep(1)
    cmd.remove('not polymer')
    cmd.remove('byres polymer & name CA')
    cmd.save(pdb_clean, state='0')
    time.sleep(1)
    cmd.reinitialize()


# class MissingEnvironmentVariable(Exception):
#     pass

def run_fpocket(
    pdb : str,
    name : str,
    chain :str,
    state : int,
    m : float,
    M : float,
    i :int,
    D : float,
    A : int,
    p : float,
    ) -> None:
    """Detects potential binding pockets in RNA structures using fpocket.

    Args:
        pdb (str): Path to input .pdb file.
        chain (str): Chain identifier for desired RNA chain (default='A').
        state (int): Structural state to analyze.
        m (float): Min. a-sphere radius in angstroms (default=3.0).
        M (float): Max. a-sphere radius in angstroms (default=5.7).
        i (int): Min. number of a-spheres per pocket (default=42).
        D (float): a-sphere clustering distance in angstroms (default=1.65).
    """
    # Prints announcement that fpocket is searching for pockets.
    name = os.path.basename(pdb)[0:-4]
    print(f'***** POCKET HUNTING {name} *****')
    # Runs fpocket bash commands
    # Try to find fpocket in the current conda/mamba/micromamba environment
    env_prefix = os.environ.get('CONDA_PREFIX', os.environ.get('MAMBA_ROOT_PREFIX', sys.prefix))
    fpocket_env_path = os.path.join(env_prefix, 'bin', 'fpocket')
    if os.path.isfile(fpocket_env_path) and os.access(fpocket_env_path, os.X_OK):
        fpocket_path = fpocket_env_path
    else:
        fpocket_path = shutil.which('fpocket')
    if not fpocket_path:
        raise FileNotFoundError('fpocket executable not found in current environment or PATH. Ensure fpocket is installed and available.')
    cmd = [fpocket_path, '-f', pdb, '-k', chain, '-l', str(state), '-m', str(m), '-M', str(M), '-i', str(i), '-D', str(D), '-A', str(A), '-p', str(p), '-w', 'p']
    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
    stdout, stderr = process.communicate()
    print(stdout)
    print(stderr)
    if 'EnvironmentLocationNotFound' in stdout or 'EnvironmentLocationNotFound' in stderr:
        raise OSError('Unable to run fpocket because the fpocketR environment does not exist.\nInstall the fpocketR environment.')


def file_fpocket(pdb : str, state : int, out : str, yes : bool) -> tuple[str, bool]:
    """Moves fpocket outputs into designated output directory.
       Default directory name specifies the fpocket parameters used.
       Manages overwriting files/directories if thet already exist.

    Args:
        pdb (str): Path to input .pdb file.
        state (int): Structural state to analyze.
        out (str): name of fpocket output parent directory name.
        y (boolean): Overwrites output files and directories with same name.

    Returns:
        str: path to output directory contianing fpocket outputs for analysis
    """
    # Moves fpocket output directories into a shared directory.
    source_dir = f'{pdb.rsplit(".")[0]}_out'
        
    if not os.path.isdir(source_dir):
        raise FileNotFoundError(f'fpocket output directory does not exist: {source_dir}.')

    if state is None:
        dest_dir = os.path.join(out, f'{pdb.rsplit(".")[0]}_out')
    else:
        dest_dir = os.path.join(out, f'{pdb.rsplit(".")[0]}_state{state}_out')

    if source_dir == dest_dir:
        analysis = source_dir

    elif not os.path.isdir(dest_dir):
        analysis = shutil.move(source_dir, dest_dir)

    # Prompts user to overwrite an existing file with same name.
    elif os.path.isdir(dest_dir):
        if not yes:
            remove = input(
                'A directory already exists with this name.\n'
                f'{dest_dir}\n\n'
                'Overwrite directory? [y/n]: '
            )
            print()
        if yes or remove in ('y', 'Y', 'yes', 'Yes'):
            yes = True
            shutil.rmtree(dest_dir)
            analysis = shutil.move(source_dir, dest_dir)
        else:
            print(
                'Exiting program. \n'
                'The name of the output directory can '
                'be changed with the --name flag.'
            )
            exit()

    # Adds state identifier number to file names.
    if state is not None:
        for file in os.listdir(dest_dir):
            file = os.path.join(dest_dir, file)
            if not os.path.isfile(file):
                continue

            head_tail = os.path.split(file)
            filename_ext = head_tail[1].split(pdb[0:-4])
            src = os.path.join(file)
            dst = os.path.join(
                head_tail[0], f'{pdb[0:-4]}_state{state}{filename_ext[1]}')

            # check if the file doesn't exist
            if not os.path.exists(dst):
                os.rename(src, dst)

    return analysis, yes
