""" This module finds the electrostatics from a 3D potential map in Gaussian
    (Software) cube format at the points on the surface of protein. The points
    on the surface of the protein are obtained by MSMS program from MGL Tools

    Author: Dibyajyoti Maity

Usage:  md_davis electrostatics [options] PDB_FILE OUTPUT_DIRECTORY

Options:
  -m, --msms PATH                   full path to MSMS executable [default: ~/]
  -d, --delphi PATH                 full path to Delphi executable [default: ~/]
  -r, --radius PATH                 path to radius file
  -c, --charge PATH                 path to charge file
  -g, --grid_size <odd_int>         Grid size to use for Delphi calculation
  -v, --vertices <filename.vert>    vertices file generated by MSMS
  -s, --surface <surface.pdb>       file containing surface vertices created by
                                    running 'md_davis vert2pdb' on vertex file
                                    created by MSMS
  --surf               Whether to calculate the electrostatic
                                    potential on the surface or not
  --center                          Center the grid for Delphi at the origin
"""

import glob
import os
import click
import subprocess
import pandas
import warnings


def dir_path(path):
    if os.path.isdir(path):
        return path
    else:
        raise NotADirectoryError(path)


# TODO: Remove hard coded application paths and names
def run_msms(pdb_file, output_directory=None, msms=None):
    """ Calculate triangulated surface using MSMS program """
    basename, _ = os.path.splitext(os.path.basename(pdb_file))

    if output_directory is None:
        output_directory = './'
    elif output_directory[-1] != '/':
        output_directory = output_directory + '/'
    else:
        pass

    if not os.path.isdir(output_directory):
        os.mkdir(output_directory)

    msms = os.path.abspath(msms)
    # TODO: Write a function to convert PDB to xyzrn
    msms_path, msms_executable = os.path.split(msms)

    with open(output_directory + basename + '.xyz', 'w') as xyz_file:
        subprocess.run([msms_path + '/pdb_to_xyzrn', os.path.abspath(pdb_file)],
                       stdout=xyz_file, cwd=msms_path)

    # TODO: MSMS is not free for commercial use, replace in the future
    # and also address the discrepancy between surface from MSMS and the
    # surface used by Delphi
    subprocess.run([msms, "-if", basename + ".xyz", "-of", basename,
                    "-probe_radius", "1.4"], cwd=output_directory)
    return output_directory + basename + '.vert'


def run_delphi(pdb_file, output_directory, output_filename,
    delphi_path, radius_file=None, charge_file=None, pqr=False, grid_size=101, surface=None, center=False):
    """ Run Delphi on protein surface created by MSMS program """
    # TODO: Rewrite using template string
    if not os.path.isdir(output_directory):
        os.mkdir(output_directory)

    parameters = [
        f'gsize={grid_size}',
        f'salt=0.15',
        f'exdi=80',
        f'linit=2000',
        f'maxc=0.0000000001',
        f'out(phi, file="{output_directory}/{output_filename}.cube", format="cube")',
    ]

    if pqr:
        parameters += [f'in(pdb,file="{pdb_file}")']
    else:
        this_script_path = os.path.dirname(os.path.realpath(__file__))
        if radius_file is None:
            radius_file = this_script_path + '/charmm.siz'

        if charge_file is None:
            charge_file = this_script_path + '/charmm.crg'

        parameters += [
            f'in(pdb,file="{pdb_file}")',
            f'in(siz,file="{radius_file}")',
            f'in(crg,file="{charge_file}")'
        ]

    if center:
        parameters += ['acenter(0,0,0)']
    if surface:
        parameters += [
            f'in(frc,file="{surface}")',
            f'out(frc, file="{output_directory}/{output_filename}.pot")',
            f'site(Atom, Potential, Reaction, Coulomb, Field)',
        ]
    print('\n'.join(parameters) + '\n', file=open(f'{output_filename}_tmp.prm', 'w'))
    subprocess.run([delphi_path, f'{output_filename}_tmp.prm'])
    os.remove(f'{output_filename}_tmp.prm')


def parse_electrostatic_potential(potential_file):
    df = pandas.read_fwf(potential_file, skiprows=12, skipfooter=2,
                         dtype={'resSeq': int}, engine='python',
                         names=['name', 'resName', 'chainID', 'resSeq',
                                'potential', 'reaction', 'coulomb',
                                'Ex', 'Ey', 'Ez'
                         ],
                         widths=[5, 3, 3, 9, 10, 10, 10, 10, 10, 10]
    )

    if df['chainID'].isnull().values.any():
        warnings.warn('Chain ID missing. '
                      'Results may be incorrect for multichain proteins. '
                      'Check if the structure had proper chain IDs')

    df['chainID'].fillna('A', inplace=True)
    output = {}
    chain = 0
    for _, data in df.groupby(['chainID'], as_index=False):
        grouped_df = data.groupby(['resSeq'], as_index=False)['potential']
        potential = grouped_df.sum()
        potential.rename(columns={'potential':'total'}, inplace=True)
        potential['mean'] = grouped_df.mean()['potential']
        output[f'chain {chain}'] = potential
        chain += 1
    return output


CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])


@click.command(name='electrostatics', context_settings=CONTEXT_SETTINGS)
@click.option('-m', '--msms', metavar='PATH', type=click.Path(exists=True),
              required=True, help='full path to MSMS executable')
@click.option('-d', '--delphi', metavar='PATH', type=click.Path(exists=True),
              required=True, help='full path to Delphi executable')
@click.option('-r', '--radius', metavar='PATH', type=click.Path(exists=True),
              help='path to radius file')
@click.option('-c', '--charge', metavar='PATH', type=click.Path(exists=True),
              help='path to charge file')
@click.option('-g', '--grid_size', metavar='<odd_int>', default=101,
              help='Grid size to use for Delphi calculation')
@click.option('--surface', is_flag=True,
              help='Calculate electrostatic potential on the surface')
@click.option('--center', help='Center the grid for Delphi at the origin')
@click.option('-o','--output', metavar='PATH', type=click.Path(),
              help='Directory to place the output files.')
@click.argument('pdb_files', nargs=-1)
def main(pdb_files, output, msms, delphi, radius, charge,
         grid_size, surface, center):
    """ Get the electrostatic potential on the surface points generated by MSMS """

    this_script_path = os.path.dirname(os.path.realpath(__file__))

    for pdb_file in pdb_files:
        output_filename = os.path.splitext(os.path.basename(pdb_file))[0]

        surface_file = None
        if surface:
            vert_file = run_msms(pdb_file=pdb_file,
                                 output_directory=output,
                                 msms=msms)
            surface_file = f"{output}/{output_filename}_surf.pdb"
            # TODO: remove this subprocess run with a function
            subprocess.run(['python', f'{this_script_path}/vert2pdb.py',
                            vert_file, pdb_file, '-o', surface_file])

        run_delphi(pdb_file=pdb_file,
                output_directory=output,
                output_filename=output_filename,
                surface=surface_file,
                delphi_path=delphi,
                radius_file=radius,
                charge_file=charge,
                grid_size=grid_size,
                center=center,
        )


if __name__ == '__main__':
    main()
