#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# CoSolvKit
#
# Utils functions
#

import contextlib
import os
import tempfile
import shutil
import subprocess
import importlib
from collections import defaultdict

import numpy as np
import mdtraj
from MDAnalysis import Universe
from openmm.app import *
from openmm import *
import pdbfixer

class MutuallyExclusiveParametersError(Exception):
    """A custom exception.

    :param Exception: this is a custom exception for mutually exclusive parameters
    :type Exception: Exception
    """
    pass   


# def add_harmonic_restraints(prmtop, inpcrd, system, atom_selection, k=1.0):
#     """Add harmonic restraints to the system.

#     Args:
#         prmtop (AmberPrmtopFile): Amber topology object
#         inpcrd (AmberInpcrdFile): Amber coordinates object
#         system (System): OpenMM system object created from the Amber topology object
#         atom_selection (str): Atom selection (see MDTraj documentation)
#         k (float): harmonic force restraints in kcal/mol/A**2 (default: 1 kcal/mol/A**2)

#     Returns:
#         (list, int): list of all the atom ids on which am harmonic force is applied, index with the System of the force that was added

#     """
#     mdtop = mdtraj.Topology.from_openmm(prmtop.topology)
#     atom_idxs = mdtop.select(atom_selection)
#     positions = inpcrd.positions

#     # Tranform constant to the right unit
#     k = k * unit.kilocalories_per_mole / unit.angstroms**2
#     k = k.value_in_unit_system(unit.md_unit_system)

#     if atom_idxs.size == 0:
#         print("Warning: no atoms selected using: %s" % atom_selection)
#         return ([], None)

#     # Take into accoun the periodic condition
#     # http://docs.openmm.org/latest/api-python/generated/simtk.openmm.openmm.CustomExternalForce.html
#     force = CustomExternalForce("k * periodicdistance(x, y, z, x0, y0, z0)^2")
    
#     harmonic_force_idx = system.addForce(force)
    
#     force.addGlobalParameter("k", k)
#     print(force.getGlobalParameterDefaultValue(0))
#     force.addPerParticleParameter("x0")
#     force.addPerParticleParameter("y0")
#     force.addPerParticleParameter("z0")
    
#     for atom_idx in atom_idxs:
#         #print(atom_idx, positions[atom_idx].value_in_unit_system(units.md_unit_system))
#         force.addParticle(int(atom_idx), positions[atom_idx].value_in_unit_system(unit.md_unit_system))

#     return atom_idxs


# def update_harmonic_restraints(simulation, k=1.0):
#     """Update harmonic restraints force
    
#     Args:
#         simulation (Simulation): OpenMM simulation object
#         k (float): new harmonic force constraint in kcal/mol/A**2 (default: 1 kcal/mol/A**2)

#     """
#     # Tranform constant to the right unit
#     k = k * unit.kilocalories_per_mole / unit.angstroms**2
#     k = k.value_in_unit_system(unit.md_unit_system)
    
#     simulation.context.setParameter('k', k)

def fix_pdb(pdbfile: str, pdbxfile: str, keep_heterogens: bool=False) -> tuple[Topology, list]:
    """Fixes common problems in PDB such as:
            - missing atoms
            - missing residues
            - missing hydrogens
            - remove nonstandard residues

    :param pdbfile: pdb string old format
    :type pdbfile: str
    :param pdbxfile: pdb string new format
    :type pdbxfile: str
    :param keep_heterogens: if False all heterogen atoms but waters are deleted, defaults to False
    :type keep_heterogens: bool, optional
    :return: new topology and positions
    :rtype: tuple[Topology, list]
    """
    fixer = pdbfixer.PDBFixer(pdbfile=pdbfile, pdbxfile=pdbxfile)
    fixer.findMissingResidues()
    
    chains = list(fixer.topology.chains())
    keys = fixer.missingResidues.keys()
    for key in list(keys):
        chain = chains[key[0]]
        if key[1] == 0 or key[1] == len(list(chain.residues())):
            del fixer.missingResidues[key]

    if not keep_heterogens:
        fixer.removeHeterogens(keepWater=True)

    fixer.findMissingAtoms() 
    fixer.addMissingAtoms()
    fixer.addMissingHydrogens(7)
    return fixer.topology, fixer.positions
    
def add_variants(topology: Topology, positions: list, variants: list=list()) -> tuple[Topology, list]:
    """Adds variants for specific protonation states.

    :param topology: openmm topology
    :type topology: Topology
    :param positions: openmm positions
    :type positions: list
    :param variants: list of variants to apply for the protonation states, defaults to list()
    :type variants: list, optional
    :return: topology and positions with added protonation states
    :rtype: tuple[Topology, list]
    """
    modeller = Modeller(topology, positions)
    added_variants = modeller.addHydrogens(variants=variants)
    return modeller.topology, modeller.positions
