import logging
import os
import subprocess

from dataclasses import dataclass
from textwrap import dedent

from rdkit import Chem
from rdkit.Chem import AllChem

from autodft.utils.autodft_utils import charge_from_smiles, multiplicity_from_smiles


logger = logging.getLogger(__name__)


# CONSTANTS
CREST_SCRIPT = 'run_crest.sh'


@dataclass(kw_only=True)
class CrestJob:
    """A class that can create, submit, and process a CREST job from
    a .xyz file.

    Creating an instance of this class leads to 4 processes:
        1. Creation of a shell script (.sh) to submit the job to a SLURM cluster
        2. Submission of the job to the computing cluster
        3. Organization of the resulting files
        4. Extraction of cartesian coordinates (as text blocks) of the
            generated conformers, which are stored in the cartesians attribute
            as a list of strings

    Example usage:
        from crest import CrestJob
        crestjob = CrestJob(molname='test', xyzfile='test.xyz', e_cutoff=2)
        for conf_cartesians in crestjob.cartesians:
            print(conf_cartesians)
    """

    molname: str
    xyzfile: str
    charge: int
    multiplicity: int
    e_cutoff: float
    
    method: str = '--gfn2'
    solvent: str = ''
    nprocshared: int = 16
    mem: str = '12gb'
    time: str = '6:00:00'
    keywords: str = None
    
    rm_extra_files: bool = True
    rm_xyzfile: bool = False
    

    # Defined in __post_init__
    conf_cartesians: list[str] = None

    def __post_init__(self) -> None:
        """Perform CREST calculation"""

        self.write_sh_script()
        self.run()
        self.cleanup()
        self.conf_cartesians = self.get_conf_cartesians()

    def write_sh_script(self) -> None:
        """Writes a shell script to submit a CREST job for the provided
        input file"""

        crest_command = (f'crest {self.xyzfile} --{self.method}'
                         f' --chrg {self.charge} --uhf {self.multiplicity-1}'
                         f' --ewin {self.e_cutoff} -T {self.nprocshared}')
        if self.solvent:
            crest_command += f' -g {self.solvent}'
        if self.keywords is not None:
            crest_command += f' {self.keywords}'

        with open(CREST_SCRIPT, 'w') as f:
            f.write(dedent(f"""\
                           #!/bin/sh
                           
                           #SBATCH --nodes=1
                           #SBATCH --ntasks=1
                           #SBATCH --cpus-per-task={self.nprocshared}
                           #SBATCH --mem={self.mem}
                           #SBATCH --time={self.time}
                           #SBATCH --output={self.molname}_crest.out
                           #SBATCH --job-name={self.molname[-8:]}
                           
                           {crest_command}
                           """))

    def run(self) -> None:
        """Submit a CREST job to the computing cluster"""

        logger.info('Submitting CREST job to SLURM...')
        subprocess.run(['sbatch', '-W', CREST_SCRIPT])
        

    def cleanup(self) -> None:
        """Organizes files after completion of the CREST job"""

        # Remove unnecessary files
        if self.rm_extra_files:
            to_keep = [f'{self.molname}.out',
                       f'{self.molname}_crest.out',
                       'crest_conformers.xyz',
                       f'autodft_{self.molname}.json',
                       'autodft_parameters.yaml']
            if not self.rm_xyzfile:
                to_keep.append(self.xyzfile)
            to_remove = [f for f in os.listdir() if f not in to_keep]
            for f in to_remove:
                try:
                    os.remove(f)
                except OSError:
                    print(f'Unable to delete: {f}')

        # Rename output file
        os.rename('crest_conformers.xyz',
                  f'{self.molname}_crest_conformers.xyz')

        # Check for normal termination of CREST job
        with open(f'{self.molname}_crest.out') as f:
            last_line = f.readlines()[-1]
            if 'CREST terminated normally.' not in last_line:
                raise CRESTAbnormalTerminationError()
        logger.info('CREST terminated normally.')

    def get_conf_cartesians(self) -> None:
        """Returns cartesian coordinates for each conformer generated by CREST.
        The conformers are ordered starting with the lowest energy."""

        xyz_blocks = []

        with open(self.molname + '_crest_conformers.xyz') as f:
            n_atoms = int(f.readline().strip())
            xyz_block_length = n_atoms + 2

            f.seek(0)
            data = f.read().splitlines(keepends=True)
            xyz_blocks = [data[i:i+xyz_block_length]
                          for i in range(0, len(data), xyz_block_length)]

        conf_cartesians = [''.join(xyz_lines[2:]) for xyz_lines in xyz_blocks]
        return conf_cartesians


@dataclass(kw_only=True)
class CrestJobFromSmiles:

    molname: str
    smiles: str
    e_cutoff: float
    
    method: str = '--gfn2'
    solvent: str = ''
    nprocshared: int = 16
    mem: str = '12gb'
    time: str = '6:00:00'
    keywords: str = None
    
    rm_extra_files: bool = True

    # Defined in __post_init__
    charge: int = None
    multiplicity: int = None
    conf_cartesians: list[str] = None

    def __post_init__(self) -> None:
        """Perform CREST calculation"""

        self.smiles = Chem.CanonSmiles(self.smiles)
        self.charge = charge_from_smiles(self.smiles)
        self.multiplicity = multiplicity_from_smiles(self.smiles)

        # Run CREST
        self.make_xyzfile()
        xyzfile = f'{self.molname}.xyz'
        args = {i: j for i, j in vars(self).items()
                if i not in ['smiles', 'conf_cartesians']}
        crest_job = CrestJob(xyzfile=xyzfile, rm_xyzfile=True, **args)

        # Get and store cartesian coordinates
        self.conf_cartesians = crest_job.conf_cartesians

    def make_xyzfile(self) -> None:
        """Generates an .xyz file (CREST input) from a SMILES string"""

        xyz_name = f'{self.molname}.xyz'

        # Embed the molecule to generate an initial 3D structure
        logger.info(f'Generating .xyz file (input for CREST): {xyz_name}')
        mol = Chem.MolFromSmiles(self.smiles)
        mol = Chem.AddHs(mol)
        AllChem.EmbedMolecule(mol, randomSeed=0xf00d,
                              useExpTorsionAnglePrefs=True,
                              useBasicKnowledge=True)
        if mol.GetNumConformers() == 0:
            raise Struct3DFromSMILESError()

        # Minimization with UFF force field
        while True:
            converged = not AllChem.UFFOptimizeMolecule(mol, maxIters=200)
            if converged:
                break

        # Write to file
        Chem.MolToXYZFile(mol, xyz_name)


class Struct3DFromSMILESError(Exception):
    """Raised when 3D generation from a SMILES string failed.
    Check stereochemical specification for highly bridged systems."""
    pass


class CRESTAbnormalTerminationError(Exception):
    """Raised when a CREST job does not terminate normally"""
    pass
