# Objects for making cube files.

from pathlib import Path
import subprocess
import tempfile
import asyncio
import os

try:
    from pyscf.tools import cubegen

except ModuleNotFoundError:
    # No PySCF.
    cubegen = None

from digichem.exception.base import File_maker_exception
from digichem.file import File_converter
import digichem.file.types as file_types
import digichem.log
from digichem.file.base import File_maker, Dummy_file_maker
from digichem.memory import Memory
from digichem.misc.io import expand_path


def sanitize_modern_cubes(cube_file_path):
    """
    'Sanitize' or 'fix' modern cube files to be compatible with older parsers.
    
    Modern cubes can store more than one density per coordinate, but many programs don't like this.
    
    :param cube_file: Path to the cube file to modify.
    """
    cube_file_path = Path(cube_file_path)
    with open(cube_file_path, "rt") as old_cube_file, \
        tempfile.NamedTemporaryFile("wt", dir = cube_file_path.parent, delete = False) as temp_cube:
        for line_index, line in enumerate(old_cube_file):
            # Line 3 (index 2) contains this header:
            #   -11  -10.884822  -13.392441  -10.884822    1
            # The final '1' is unsupported on older software.
            if line_index == 2:
                line_split = line.split()
                if len(line_split) == 5:
                    # We have a line that needs cleaning.
                    if line_split[-1] != "1":
                        digichem.log.get_logger().warning(
                            "Cube file {} contains more than 1 density ('{}' densities found), this file may no longer be fully compliant".format(
                                cube_file_path,
                                line_split[-1]
                            )
                        )
                    
                    line = " " + " ".join(line_split[:-1]) + "\n"
            
            temp_cube.write(line)
            
        # Replace old cube with new cube.
        os.rename(temp_cube.name, cube_file_path)
    

class Fchk_to_cube(File_converter):
    """
    Class for handling file references to Gaussian cube files.
    
    We currently rely on Gaussian's cubegen facility to make cube files. There are a number of non-trivial problems with this approach:
     - Cubegen is part of the (very expensive) Gaussian package and so is not readily available on most machines used for post-processing.
     - The cubegen interface is poor, which makes effective error reporting particularly challenging.
     - Cubegen only takes .fchk files as input, while the rest of this analysis package reads the .log file.
       - .fchk files are not even generated by Gaussian automatically, they must be created from .chk files using another buggy Gaussian program: formchk
       - The requirement on these .fchk files is unnecessary as the data is also available in the .log file, but cubegen cannot read it.
    """
    
    # Text description of our input file type, used for error messages etc.
    #input_file_type = "fchk"
    input_file_type = file_types.gaussian_fchk_file
    # Text description of our output file type, used for error messages etc.
    output_file_type = file_types.gaussian_cube_file
    
    def __init__(
            self,
            *args,
            fchk_file = None,
            cubegen_type = "MO",
            orbital = "HOMO",
            npts = 0,
            cube_file = None,
            num_cpu = None,
            memory = None,
            cubegen_executable = "cubegen",
            sanitize = False,
            multithreading = None,
            **kwargs):
        """
        Constructor for Fchk_to_cube objects.
        
        See Image_maker for a full signature.
        
        :param output: The filename/path to the cube file (this path doesn't need to point to a real file yet; we will use this path to write to).
        :param fchk_file: Optional fchk_file to use to generate this cube file.
        :param cubegen_type: The type of orbital that will be included in the cube file when we make it. Possible values are MO, AMO and BMO (for MO, alpha MO and beta MO respectively).
        :param orbital: The orbital to be included in the cube file when we make it. Possible values are 'HOMO', 'LUMO' or the integer level of the desired orbital.
        :param npts: The 'npts' option of cubegen, controls how detailed the resulting file is. Common options are 0 (default), -2 ('low' quality), -3 (medium quality), -4 (very high quality).
        :param cube_file: An optional file path to an existing cube file to use. If this is given (and points to an actual file), then a new cube will not be made and this file will be used instead.
        :param num_cpu: The number of CPUs/processes for cubegen to use. If not given, the number of CPUs available will be used.
        :param memory: The amount of memory for cubegen to use.
        :param sanitize: Whether to modify the cube file to make it compatible with older software.
        """
        super().__init__(*args, input_file = fchk_file, existing_file = cube_file, **kwargs)
        self.cubegen_type = cubegen_type
        self.orbital = orbital
        self.npts = npts
        memory = memory if memory is not None else "3 GB"
        self.memory = Memory(memory)
        # Default to number of CPUs of the system.
        # If we've not been asked to run multithreaded, use 1 CPU.
        if multithreading is None or multithreading == "pool":
            self.num_cpu = 1
        
        else:
            self.num_cpu = int(num_cpu) if num_cpu is not None else os.cpu_count()
        
        self.cubegen_executable = expand_path(cubegen_executable)
        self.sanitize = sanitize
        
    @classmethod
    def from_options(self, output, *, fchk_file = None, cubegen_type = "MO", orbital = "HOMO", num_cpu = None, options, **kwargs):
        """
        Constructor that takes a dictionary of config like options.
        """        
        return self(
            output,
            fchk_file = fchk_file,
            cubegen_type = cubegen_type,
            orbital = orbital,
            npts = options['render']['orbital']['cube_grid_size'].to_gaussian(),
            dont_modify = not options['render']['enable_rendering'],
            cubegen_executable = options['external']['cubegen'],
            sanitize = options['render']['safe_cubes'],
            num_cpu = num_cpu,
            multithreading = options['external']['cubegen_parallel'],
            **kwargs
        )
    
    def get_parallel(self, name = 'file', cpus = 1):
        """
        Generate the file represented by this object in a parallel context.

        get_parallel() will be called by a higher level function as an argument to ThreadPoolExecutor.map() or similar,
        to generate many files simultaneously. This is useful for slow operations (such as cube generation) that are
        difficult to parallelise individually, but easy to parallelise across multiple files.

        :param name: The file to generate.
        :param cpus: The number of CPUs this operation should use, nearly always 1.
        """
        # Temporarily set CPUs to 1 (in-case it got changed somewhere else.)
        old_cpus = self.num_cpu
        self.num_cpu = cpus
        try:
            self.get_file(name = name)
            
        finally:
            self.num_cpu = old_cpus

    @property
    def signature(self):
        """
        Get the call signature to run cubegen (as can be passed to Popen etc)
        """
        return [
            "{}".format(self.cubegen_executable),
            # Some versions of cubegen respect this num_threads arg, some don't.
            "{}".format(self.num_cpu),
            "{}={}".format(self.cubegen_type, self.orbital),
            str(self.input_file),
            str(self.output),
            str(self.npts)
        ]

    @property
    def env(self):
        """
        Get the environmental variables require to run cubegen (as can be passed to Popen etc).
        """
        return dict(
            os.environ,
            GAUSS_MEMDEF = str(self.memory),
            # Some versions of cubegen only respect this way of changing threads.
            OMP_NUM_THREADS = str(self.num_cpu)
        )

    async def make_files_async(self):
        """
        Make the files referenced by this object.
        """        
        try:
            cubegen_proc =  await asyncio.create_subprocess_exec(
                # Unless subprocess.run(), create_subprocess_exec() takes and expanded list as first arg...
                *self.signature,
                # Capture both stdout and stderr.
                # It appears that cubegen writes everything (including error messages) to stdout, but it does return meaningful exit codes.
                stdout = subprocess.PIPE,
                stderr = subprocess.STDOUT,
                env = self.env
            )
        except FileNotFoundError:
            raise File_maker_exception(self, "Could not locate cubegen executable '{}'".format(self.cubegen_executable))
        
        stdout = ""
        stderr = ""
        # Wait for completion.
        while cubegen_proc.returncode is None:
            retval = await cubegen_proc.communicate()
            stdout += retval[0].decode()
            stderr += retval[1].decode()

        # If something went wrong, dump output.
        if cubegen_proc.returncode != 0:
            # An error occured.
            # Check if the input file exists, if not this is probably what went wrong.
            message = "Cubegen did not exit successfully"
            if not Path(str(self.input_file)).exists():
                message += " (probably because the input file '{}' could not be found)".format(self.input_file)
            message += ":\n{}".format(stdout)
            raise File_maker_exception(self, message)
        else:
            # Everything appeared to go ok.
            # Dump cubegen output if we're in debug.
            digichem.log.get_logger().debug(stdout)
        
        if self.sanitize:
            sanitize_modern_cubes(self.output)

            
    def make_files(self):
        """
        Make the files referenced by this object.
        """
        try:
            cubegen_proc =  subprocess.run(
                self.signature,
                # Capture both stdout and stderr.
                # It appears that cubegen writes everything (including error messages) to stdout, but it does return meaningful exit codes.
                stdout = subprocess.PIPE,
                stderr = subprocess.STDOUT,
                universal_newlines = True,
                env = self.env
            )
        except FileNotFoundError:
            raise File_maker_exception(self, "Could not locate cubegen executable '{}'".format(self.cubegen_executable))
        
        # If something went wrong, dump output.
        if cubegen_proc.returncode != 0:
            # An error occured.
            # Check if the input file exists, if not this is probably what went wrong.
            message = "Cubegen did not exit successfully"
            if not Path(str(self.input_file)).exists():
                message += " (probably because the input file '{}' could not be found)".format(self.input_file)
            message += ":\n{}".format(cubegen_proc.stdout)
            raise File_maker_exception(self, message)
        else:
            # Everything appeared to go ok.
            # Dump cubegen output if we're in debug.
            digichem.log.get_logger().debug(cubegen_proc.stdout)
        
        if self.sanitize:
            sanitize_modern_cubes(self.output)


class Fchk_to_spin_cube(Fchk_to_cube):
    """
    A variation of the cube maker designed for making spin density cubes.
    
    Note that this class exists mainly for convenience; either class can be used to generate cubes of any supported type.
    """
    
    def __init__(self, *args, spin_density = "SCF", **kwargs):
        """
        Constructor for Fchk_to_cube objects.
        
        See Fchk_to_cube for a full signature.
        
        :param output: The filename/path to the cube file (this path doesn't need to point to a real file yet; we will use this path to write to).
        :param fchk_file: Optional fchk_file to use to generate this cube file.
        :param spin_density: The density to use; you almost certainly want SCF unless perhaps you are using MP2...
        :param npts: The 'npts' option of cubegen, controls how detailed the resulting file is. Common options are 0 (default), -2 ('low' quality), -3 (medium quality), -4 (very high quality).
        :param cube_file: An optional file path to an existing cube file to use. If this is given (and points to an actual file), then a new cube will not be made and this file will be used instead.
        """
        return super().__init__(*args, cubegen_type = "Spin", orbital = spin_density, **kwargs)
    
    @classmethod
    def from_options(self, output, *, fchk_file = None, spin_density = "SCF", num_cpu = None, options, **kwargs):
        """
        Constructor that takes a dictionary of config like options.
        """        
        return self(
            output,
            fchk_file = fchk_file,
            spin_density = spin_density,
            npts = options['render']['spin']['cube_grid_size'].to_gaussian(),
            dont_modify = not options['render']['enable_rendering'],
            cubegen_executable = options['external']['cubegen'],
            sanitize = options['render']['safe_cubes'],
            num_cpu = num_cpu,
            multithreading = options['external']['cubegen_parallel'],
            **kwargs
        )

class Fchk_to_density_cube(Fchk_to_cube):
    """
    A variation of the cube maker designed for making density cubes.
    
    Note that this class exists mainly for convenience; either class can be used to generate cubes of any supported type.
    """
    
    def __init__(self, *args, density_type = "SCF", **kwargs):
        """
        Constructor for Fchk_to_cube objects.
        
        See Fchk_to_cube for a full signature.
        
        :param output: The filename/path to the cube file (this path doesn't need to point to a real file yet; we will use this path to write to).
        :param fchk_file: Optional fchk_file to use to generate this cube file.
        :param density_type: The density to use.
        :param npts: The 'npts' option of cubegen, controls how detailed the resulting file is. Common options are 0 (default), -2 ('low' quality), -3 (medium quality), -4 (very high quality).
        :param cube_file: An optional file path to an existing cube file to use. If this is given (and points to an actual file), then a new cube will not be made and this file will be used instead.
        """
        super().__init__(*args, cubegen_type = "Density", orbital = density_type, **kwargs)
        self.type = density_type
    
    @classmethod
    def from_options(self, output, *, fchk_file = None, density_type = "SCF", num_cpu = None, options, **kwargs):
        """
        Constructor that takes a dictionary of config like options.
        """        
        return self(
            output,
            fchk_file = fchk_file,
            density_type = density_type,
            npts = options['render']['density']['cube_grid_size'].to_gaussian(),
            dont_modify = not options['render']['enable_rendering'],
            cubegen_executable = options['external']['cubegen'],
            sanitize = options['render']['safe_cubes'],
            num_cpu = num_cpu,
            multithreading = options['external']['cubegen_parallel'],
            **kwargs
        )

class Fchk_to_nto_cube(Fchk_to_cube):
    """
    A variation of the cube maker designed for making NTO cubes.
    
    Note that this class exists mainly for convenience; either class can be used to generate cubes of any supported type.
    """
    
    def __init__(self, *args, orbital = "HOMO", **kwargs):
        """
        Constructor for Fchk_to_cube objects.
        
        See Fchk_to_cube for a full signature.
        
        :param output: The filename/path to the cube file (this path doesn't need to point to a real file yet; we will use this path to write to).
        :param fchk_file: Optional fchk_file to use to generate this cube file.
        :param orbital: The NTO to plot.
        :param npts: The 'npts' option of cubegen, controls how detailed the resulting file is. Common options are 0 (default), -2 ('low' quality), -3 (medium quality), -4 (very high quality).
        :param cube_file: An optional file path to an existing cube file to use. If this is given (and points to an actual file), then a new cube will not be made and this file will be used instead.
        """
        super().__init__(*args, cubegen_type = "MO", orbital = orbital, **kwargs)
    
    @classmethod
    def from_options(self, output, *, fchk_file = None, orbital = "HOMO", num_cpu = None, options, **kwargs):
        """
        Constructor that takes a dictionary of config like options.
        """        
        return self(
            output,
            fchk_file = fchk_file,
            orbital = orbital,
            npts = options['render']['natural_transition_orbital']['cube_grid_size'].to_gaussian(),
            dont_modify = not options['render']['enable_rendering'],
            cubegen_executable = options['external']['cubegen'],
            sanitize = options['render']['safe_cubes'],
            num_cpu = num_cpu,
            multithreading = options['external']['cubegen_parallel'],
            **kwargs
        )

class PySCF_to_cube(File_maker):
    """
    Generate cubes from a completed PySCF calculation result object.
    """

    def __init__(
            self,
            *args,
            mol,
            target,
            target_type = "density",
            npts = 80,
            cube_file = None,
            sanitize = False,
            **kwargs):
        """
        Constructor for Fchk_to_cube objects.
        
        See Image_maker for a full signature.
        
        :param output: The filename/path to the cube file (this path doesn't need to point to a real file yet; we will use this path to write to).
        :param npts: The number of points per side of the cube.
        :param cube_file: An optional file path to an existing cube file to use. If this is given (and points to an actual file), then a new cube will not be made and this file will be used instead.
        :param sanitize: Whether to modify the cube file to make it compatible with older software.
        """
        super().__init__(*args, existing_file = cube_file, **kwargs)
        self.npts = npts
        self.sanitize = sanitize
        self.target_type = target_type
        self.target = target
        self.mol = mol
        # TODO: Add some intelligence to this...
        self.type = "SCF"
        
    @classmethod
    def from_options(self, output, *, options, **kwargs):
        """
        Constructor that takes a dictionary of config like options.
        """        
        return self(
            output,
            npts = options['render']['orbital']['cube_grid_size'].translate("points"),
            dont_modify = not options['render']['enable_rendering'],
            sanitize = options['render']['safe_cubes'],
            **kwargs
        )
    
    def check_can_make(self):
        super().check_can_make()

        if cubegen is None:
            raise File_maker_exception(self, "PySCF is not available")
    
    def make_files(self):
        """
        Make the files referenced by this object.
        """
        if self.target_type == "density":
            cubegen.density(self.mol, self.output, self.target.make_rdm1(), nx=self.npts, ny=self.npts, nz=self.npts)

        elif self.target_type == "orbital":
            cubegen.orbital(self.mol, self.output, self.target, nx=self.npts, ny=self.npts, nz=self.npts)
        
        else:
            raise ValueError("Unrecognised 'target_type': {}".format(self.target_type))
        
        if self.sanitize:
            sanitize_modern_cubes(self.output)