"""A wrapper for running generic lammps scripts with injected variables"""
import os
from typing import Tuple
from matplotlib.figure import Figure
from matplotlib.axes import Axes

from phonopy import Phonopy
from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections

from my_phonolammps._phonopy_overrides import (write_hdf5_band_structure,
                                               my_phonopy_load,
                                               MyPhonopy)
from my_phonolammps._phonolammps import MyPhonolammps

from my_phonolammps.util import _print
from my_phonolammps.phonon_outputs import PhononOutputFiles
from my_phonolammps.lammps_runner import LammpsRunner, LammpsVarLine, VarType

_ID_MATRIX = [[1, 0, 0],
              [0, 1, 0],
              [0, 0, 1]]


class PhononRunner:
    """
    Runs phonolammps and phonopy calculations on a given nanowire structure
    and calculates the Γ->Χ phonon dispersion along the wire axis by default.
    Adjust `phonon_path` and `phonon_mesh` accordingly to make changes.

    The main output is an HDF5 file (default "band.hdf5").
    With the default settings, this can easily be several GB in size for a
    large (100-1000 atom) unit cell.
    """

    # pylint: disable=too-many-instance-attributes

    def __init__(self, wire_datafile: str, potential_file: str,
                 output_dir='.',
                 generic_relax='./generic_relax.lammps',
                 generic_in='./generic_in.lammps',
                 marker=None,
                 clean_up=False,
                 trash_counter_max=50):

        self.clean_up = clean_up
        self._trash_counter_max = trash_counter_max

        # nothing works if these don't exist
        generic_relax = os.path.expanduser(generic_relax)
        if not os.path.isfile(generic_relax):
            raise FileNotFoundError(f"file not found '{generic_relax}'")
        self.generic_relax = generic_relax

        generic_in = os.path.expanduser(generic_in)
        if not os.path.isfile(generic_in):
            raise FileNotFoundError(f"file not found '{generic_in}'")
        self.generic_in = generic_in

        potential_file = os.path.expanduser(potential_file)
        if not os.path.isfile(potential_file):
            raise FileNotFoundError(f"file not found '{potential_file}'")
        self.potential_file = potential_file

        # validate paths before assigning
        wire_datafile = os.path.expanduser(wire_datafile)
        if os.path.isfile(wire_datafile):
            self.wire_datafile = wire_datafile
        else:
            raise FileNotFoundError(f"file not found '{wire_datafile}'")
        self.outputs = PhononOutputFiles(output_dir, marker=marker)

        self._phonon_path = [[[0., 0., 0.], [0., 0., .5]]]
        self._phonon_mesh = [1, 1, 20]
        self._nqpoints = 51
        self._phonon = None
        self._phl = None

    def __repr__(self):
        return f"""\
PhononRunner(wire_datafile='{self.wire_datafile}',
             potential_file='{self.potential_file}',
             output_dir='{self.outputs.output_dir}',
             generic_relax='{self.generic_relax}',
             generic_in='{self.generic_in}',
             marker='{self.outputs.marker}',
             clean_up={self.clean_up},
             trash_counter_max={self._trash_counter_max})
             
             self.phonon_path={self.phonon_path if self.phonon_path else None}
             self.phonon_mesh={self.phonon_mesh if self.phonon_mesh else None}
             self.nqpoints={self.nqpoints if self.nqpoints else None}
             self.phl={self._phl if self._phl else None}
             self.outputs={self.outputs}
            """

    def __del__(self):
        if self.clean_up:
            self.outputs.clean_up()

    @property
    def phonon(self) -> MyPhonopy:
        if self._phonon is None:
            _print("loading phonon")
            load_kwargs = dict(supercell_matrix=_ID_MATRIX,
                               unitcell_filename=self.outputs.unitcell_filename,
                               force_constants_filename=self.outputs.force_constants_filename)
            self._phonon = my_phonopy_load(**load_kwargs)
        return self._phonon

    @phonon.setter
    def phonon(self, ph: Phonopy):
        if isinstance(ph, Phonopy):
            self._phonon = ph

    @property
    def phonon_path(self):
        return self._phonon_path

    @phonon_path.setter
    def phonon_path(self, _path):
        self._phonon_path = _path

    @property
    def phonon_mesh(self):
        return self._phonon_mesh

    @phonon_mesh.setter
    def phonon_mesh(self, _mesh):
        if len(_mesh) != 3:
            raise ValueError("mesh should be a 3-vector")
        self._phonon_mesh = [int(n) for n in _mesh]

    @property
    def nqpoints(self):
        return self._nqpoints

    @nqpoints.setter
    def nqpoints(self, _npoints):
        self._nqpoints = int(_npoints)

    @property
    def phl(self) -> MyPhonolammps:
        return self._phl

    def run(self,
            omit_zeros: bool = True,
            omit_zeros_thresh: float = 1e-6,
            with_eigenvectors: bool = True,
            with_group_velocities: bool = True) -> None:
        """relax structure, compute forces, and compute phonons"""
        self.relax_structure()
        self.compute_forces(omit_zeros=omit_zeros, omit_zeros_thresh=omit_zeros_thresh)
        self.compute_band_structure(with_eigenvectors=with_eigenvectors,
                                    with_group_velocities=with_group_velocities)

    def relax_structure(self, dummy: bool = False) -> None:
        """
        use LAMMPS to relax input structure
        (minimize potential energy at 0 pressure)
        """

        # aliasing for neatness
        wire = self.wire_datafile
        wire_relaxed = self.outputs.relaxed_unitcell_filename  # this file is written
        potential_file = self.potential_file

        required_vars = [
            LammpsVarLine("data_file", VarType.STRING, wire),
            LammpsVarLine("output_file", VarType.STRING, wire_relaxed),
            LammpsVarLine("potential_file", VarType.STRING, potential_file)
        ]
        lmp_runner = LammpsRunner(self.generic_relax, required_vars)
        append_commands = [
            'reset_timestep 0',
            f'dump RELAXED all atom 1 {self.outputs.relaxed_unitcell_filename}',
            'dump_modify RELAXED sort id '
            'format line "%d %d %20.15g %20.15g %20.15g"',  # IMPORTANT
            'run 0',
            'undump RELAXED'
        ]
        if dummy:
            lmp_runner.run(append_commands=append_commands,
                           skip_keywords=["minimize",
                                          "compute",
                                          "thermo_style"])
        else:
            lmp_runner.run(append_commands=append_commands)

    def compute_forces(self,
                       relaxed: bool = True,
                       omit_zeros: bool = True,
                       omit_zeros_thresh: float = 1e-6) -> None:
        """use phonolammps to compute and write the force constants"""

        # aliasing for neatness
        wire_datafile = self.wire_datafile
        relaxed_datafile = self.outputs.relaxed_unitcell_filename
        unitcell = self.outputs.unitcell_filename
        force_constants = self.outputs.force_constants_filename
        potential_file = self.potential_file

        if not relaxed:
            # re-write 'relaxed' data-file without doing minimization
            self.relax_structure(dummy=True)

        required_vars = [
            # TODO: fix these terrible names
            LammpsVarLine("data_file", VarType.STRING, wire_datafile),
            LammpsVarLine("output_file", VarType.STRING, relaxed_datafile),
            LammpsVarLine("potential_file", VarType.STRING, potential_file)
        ]

        input_file = LammpsRunner(self.generic_in,
                                  required_vars).write_final_script()

        if self._phl is None:
            self._phl = MyPhonolammps(input_file,
                                      supercell_matrix=_ID_MATRIX,
                                      show_progress=True,
                                      trash_counter_max=self._trash_counter_max)

        self._phl.write_unitcell_POSCAR(unitcell)
        self._phl.write_force_constants(force_constants,
                                        omit_zeros=omit_zeros,
                                        omit_zeros_thresh=omit_zeros_thresh)

        # delete substituted LAMMPS input file
        os.remove(input_file)

    def compute_harmonic_constants(self,
                                   relaxed: bool = True,
                                   omit_zeros: bool = True,
                                   omit_zeros_thresh: float = 1e-6) -> None:
        """use phonolammps to compute and write the force constants"""

        # aliasing for neatness
        wire_datafile = self.wire_datafile
        relaxed_datafile = self.outputs.relaxed_unitcell_filename
        hc_filename = self.outputs.harmonic_constants_filename
        potential_file = self.potential_file

        if not relaxed:
            # re-write 'relaxed' data-file without doing minimization
            self.relax_structure(dummy=True)

        required_vars = [
            LammpsVarLine("data_file", VarType.STRING, wire_datafile),
            LammpsVarLine("output_file", VarType.STRING, relaxed_datafile),
            LammpsVarLine("potential_file", VarType.STRING, potential_file)
        ]

        input_file = LammpsRunner(self.generic_in,
                                  required_vars).write_final_script()

        if self._phl is None:
            self._phl = MyPhonolammps(input_file,
                                      supercell_matrix=_ID_MATRIX,
                                      show_progress=True,
                                      trash_counter_max=self._trash_counter_max)

        self._phl.write_harmonic_constants(hc_filename,
                                           omit_zeros=omit_zeros,
                                           omit_zeros_thresh=omit_zeros_thresh)

        # delete substituted LAMMPS input file
        os.remove(input_file)

    def run_mesh(self,
                 with_eigenvectors: bool = True,
                 with_group_velocities: bool = True) -> None:
        """run the mesh on the phonon attribute"""
        _print("running mesh")
        self.phonon.run_mesh(self._phonon_mesh,
                             with_eigenvectors=with_eigenvectors,
                             with_group_velocities=with_group_velocities)

    def compute_band_structure(self,
                               with_eigenvectors: bool = True,
                               with_group_velocities: bool = True,
                               use_C_library: bool = True) -> None:
        """use phonopy to compute band structure"""
        if self.phonon.mesh is None:
            self.run_mesh(with_eigenvectors, with_group_velocities)

        _print(self.phonon.dynamical_matrix.__repr__())

        _print("getting qpoints and connections")
        qs, cons = get_band_qpoints_and_path_connections(band_paths=self._phonon_path,
                                                         npoints=self._nqpoints)
        _print("running band structure")
        self.phonon.run_band_structure(qs,
                                       path_connections=cons,
                                       with_eigenvectors=with_eigenvectors,
                                       with_group_velocities=with_group_velocities,
                                       use_C_library=use_C_library)

        # use alternate function defined below
        _print("writing band structure")
        write_hdf5_band_structure(self.phonon, paths=qs, filename=self.outputs.band)

    def compute_total_dos(self,
                          with_eigenvectors: bool = True,
                          with_group_velocities: bool = True):
        """use phonopy to compute the total DOS"""
        if self.phonon.mesh is None:
            self.run_mesh(with_eigenvectors, with_group_velocities)

        _print("running total dos")
        self.phonon.run_total_dos()

        _print("writing total dos")
        self.phonon.write_total_dos(filename=self.outputs.dos)

    def plot_bands(self) -> Tuple[Figure, Axes]:
        """run the default phonopy `plot_band_structure` method"""
        plt = self.phonon.plot_band_structure()
        return plt.gcf(), plt.gca()

"""A wrapper for running generic lammps scripts with injected variables"""
import os
from typing import Tuple
from matplotlib.figure import Figure
from matplotlib.axes import Axes

from phonopy import Phonopy
from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections

from my_phonolammps._phonopy_overrides import (write_hdf5_band_structure,
                                               my_phonopy_load,
                                               MyPhonopy)
from my_phonolammps._phonolammps import MyPhonolammps

from my_phonolammps.util import _print
from my_phonolammps.phonon_outputs import PhononOutputFiles
from my_phonolammps.lammps_runner import LammpsRunner, LammpsVarLine, VarType

_ID_MATRIX = [[1, 0, 0],
              [0, 1, 0],
              [0, 0, 1]]


class PhononRunner:
    """
    Runs phonolammps and phonopy calculations on a given nanowire structure
    and calculates the Γ->Χ phonon dispersion along the wire axis by default.
    Adjust `phonon_path` and `phonon_mesh` accordingly to make changes.

    The main output is an HDF5 file (default "band.hdf5").
    With the default settings, this can easily be several GB in size for a
    large (100-1000 atom) unit cell.
    """

    # pylint: disable=too-many-instance-attributes

    def __init__(self, wire_datafile: str, potential_file: str,
                 output_dir='.',
                 generic_relax='./generic_relax.lammps',
                 generic_in='./generic_in.lammps',
                 marker=None,
                 clean_up=False,
                 trash_counter_max=50):

        self.clean_up = clean_up
        self._trash_counter_max = trash_counter_max

        # nothing works if these don't exist
        generic_relax = os.path.expanduser(generic_relax)
        if not os.path.isfile(generic_relax):
            raise FileNotFoundError(f"file not found '{generic_relax}'")
        self.generic_relax = generic_relax

        generic_in = os.path.expanduser(generic_in)
        if not os.path.isfile(generic_in):
            raise FileNotFoundError(f"file not found '{generic_in}'")
        self.generic_in = generic_in

        potential_file = os.path.expanduser(potential_file)
        if not os.path.isfile(potential_file):
            raise FileNotFoundError(f"file not found '{potential_file}'")
        self.potential_file = potential_file

        # validate paths before assigning
        wire_datafile = os.path.expanduser(wire_datafile)
        if os.path.isfile(wire_datafile):
            self.wire_datafile = wire_datafile
        else:
            raise FileNotFoundError(f"file not found '{wire_datafile}'")
        self.outputs = PhononOutputFiles(output_dir, marker=marker)

        self._phonon_path = [[[0., 0., 0.], [0., 0., .5]]]
        self._phonon_mesh = [1, 1, 20]
        self._nqpoints = 51
        self._phonon = None
        self._phl = None

    def __repr__(self):
        return f"""\
PhononRunner(wire_datafile='{self.wire_datafile}',
             potential_file='{self.potential_file}',
             output_dir='{self.outputs.output_dir}',
             generic_relax='{self.generic_relax}',
             generic_in='{self.generic_in}',
             marker='{self.outputs.marker}',
             clean_up={self.clean_up},
             trash_counter_max={self._trash_counter_max})
             
             self.phonon_path={self.phonon_path if self.phonon_path else None}
             self.phonon_mesh={self.phonon_mesh if self.phonon_mesh else None}
             self.nqpoints={self.nqpoints if self.nqpoints else None}
             self.phl={self._phl if self._phl else None}
             self.outputs={self.outputs}
            """

    def __del__(self):
        if self.clean_up:
            self.outputs.clean_up()

    @property
    def phonon(self) -> MyPhonopy:
        if self._phonon is None:
            _print("loading phonon")
            load_kwargs = dict(supercell_matrix=_ID_MATRIX,
                               unitcell_filename=self.outputs.unitcell_filename,
                               force_constants_filename=self.outputs.force_constants_filename)
            self._phonon = my_phonopy_load(**load_kwargs)
        return self._phonon

    @phonon.setter
    def phonon(self, ph: Phonopy):
        if isinstance(ph, Phonopy):
            self._phonon = ph

    @property
    def phonon_path(self):
        return self._phonon_path

    @phonon_path.setter
    def phonon_path(self, _path):
        self._phonon_path = _path

    @property
    def phonon_mesh(self):
        return self._phonon_mesh

    @phonon_mesh.setter
    def phonon_mesh(self, _mesh):
        if len(_mesh) != 3:
            raise ValueError("mesh should be a 3-vector")
        self._phonon_mesh = [int(n) for n in _mesh]

    @property
    def nqpoints(self):
        return self._nqpoints

    @nqpoints.setter
    def nqpoints(self, _npoints):
        self._nqpoints = int(_npoints)

    @property
    def phl(self) -> MyPhonolammps:
        return self._phl

    def run(self,
            omit_zeros: bool = True,
            omit_zeros_thresh: float = 1e-6,
            with_eigenvectors: bool = True,
            with_group_velocities: bool = True) -> None:
        """relax structure, compute forces, and compute phonons"""
        self.relax_structure()
        self.compute_forces(omit_zeros=omit_zeros, omit_zeros_thresh=omit_zeros_thresh)
        self.compute_band_structure(with_eigenvectors=with_eigenvectors,
                                    with_group_velocities=with_group_velocities)

    def relax_structure(self, dummy: bool = False) -> None:
        """
        use LAMMPS to relax input structure
        (minimize potential energy at 0 pressure)
        """

        # aliasing for neatness
        wire = self.wire_datafile
        wire_relaxed = self.outputs.relaxed_unitcell_filename  # this file is written
        potential_file = self.potential_file

        required_vars = [
            LammpsVarLine("data_file", VarType.STRING, wire),
            LammpsVarLine("output_file", VarType.STRING, wire_relaxed),
            LammpsVarLine("potential_file", VarType.STRING, potential_file)
        ]
        lmp_runner = LammpsRunner(self.generic_relax, required_vars)
        append_commands = [
            'reset_timestep 0',
            f'dump RELAXED all atom 1 {self.outputs.relaxed_unitcell_filename}',
            'dump_modify RELAXED sort id '
            'format line "%d %d %20.15g %20.15g %20.15g"',  # IMPORTANT
            'run 0',
            'undump RELAXED'
        ]
        if dummy:
            lmp_runner.run(append_commands=append_commands,
                           skip_keywords=["minimize",
                                          "compute",
                                          "thermo_style"])
        else:
            lmp_runner.run(append_commands=append_commands)

    def compute_forces(self,
                       relaxed: bool = True,
                       omit_zeros: bool = True,
                       omit_zeros_thresh: float = 1e-6) -> None:
        """use phonolammps to compute and write the force constants"""

        # aliasing for neatness
        wire_datafile = self.wire_datafile
        relaxed_datafile = self.outputs.relaxed_unitcell_filename
        unitcell = self.outputs.unitcell_filename
        force_constants = self.outputs.force_constants_filename
        potential_file = self.potential_file

        if not relaxed:
            # re-write 'relaxed' data-file without doing minimization
            self.relax_structure(dummy=True)

        required_vars = [
            # TODO: fix these terrible names
            LammpsVarLine("data_file", VarType.STRING, wire_datafile),
            LammpsVarLine("output_file", VarType.STRING, relaxed_datafile),
            LammpsVarLine("potential_file", VarType.STRING, potential_file)
        ]

        input_file = LammpsRunner(self.generic_in,
                                  required_vars).write_final_script()

        if self._phl is None:
            self._phl = MyPhonolammps(input_file,
                                      supercell_matrix=_ID_MATRIX,
                                      show_progress=True,
                                      trash_counter_max=self._trash_counter_max)

        self._phl.write_unitcell_POSCAR(unitcell)
        self._phl.write_force_constants(force_constants,
                                        omit_zeros=omit_zeros,
                                        omit_zeros_thresh=omit_zeros_thresh)

        # delete substituted LAMMPS input file
        os.remove(input_file)

    def compute_harmonic_constants(self,
                                   relaxed: bool = True,
                                   omit_zeros: bool = True,
                                   omit_zeros_thresh: float = 1e-6) -> None:
        """use phonolammps to compute and write the force constants"""

        # aliasing for neatness
        wire_datafile = self.wire_datafile
        relaxed_datafile = self.outputs.relaxed_unitcell_filename
        hc_filename = self.outputs.harmonic_constants_filename
        potential_file = self.potential_file

        if not relaxed:
            # re-write 'relaxed' data-file without doing minimization
            self.relax_structure(dummy=True)

        required_vars = [
            LammpsVarLine("data_file", VarType.STRING, wire_datafile),
            LammpsVarLine("output_file", VarType.STRING, relaxed_datafile),
            LammpsVarLine("potential_file", VarType.STRING, potential_file)
        ]

        input_file = LammpsRunner(self.generic_in,
                                  required_vars).write_final_script()

        if self._phl is None:
            self._phl = MyPhonolammps(input_file,
                                      supercell_matrix=_ID_MATRIX,
                                      show_progress=True,
                                      trash_counter_max=self._trash_counter_max)

        self._phl.write_harmonic_constants(hc_filename,
                                           omit_zeros=omit_zeros,
                                           omit_zeros_thresh=omit_zeros_thresh)

        # delete substituted LAMMPS input file
        os.remove(input_file)

    def run_mesh(self,
                 with_eigenvectors: bool = True,
                 with_group_velocities: bool = True) -> None:
        """run the mesh on the phonon attribute"""
        _print("running mesh")
        self.phonon.run_mesh(self._phonon_mesh,
                             with_eigenvectors=with_eigenvectors,
                             with_group_velocities=with_group_velocities)

    def compute_band_structure(self,
                               with_eigenvectors: bool = True,
                               with_group_velocities: bool = True,
                               use_C_library: bool = True) -> None:
        """use phonopy to compute band structure"""
        if self.phonon.mesh is None:
            self.run_mesh(with_eigenvectors, with_group_velocities)

        _print(self.phonon.dynamical_matrix.__repr__())

        _print("getting qpoints and connections")
        qs, cons = get_band_qpoints_and_path_connections(band_paths=self._phonon_path,
                                                         npoints=self._nqpoints)
        _print("running band structure")
        self.phonon.run_band_structure(qs,
                                       path_connections=cons,
                                       with_eigenvectors=with_eigenvectors,
                                       with_group_velocities=with_group_velocities,
                                       use_C_library=use_C_library)

        # use alternate function defined below
        _print("writing band structure")
        write_hdf5_band_structure(self.phonon, paths=qs, filename=self.outputs.band)

    def compute_total_dos(self,
                          with_eigenvectors: bool = True,
                          with_group_velocities: bool = True):
        """use phonopy to compute the total DOS"""
        if self.phonon.mesh is None:
            self.run_mesh(with_eigenvectors, with_group_velocities)

        _print("running total dos")
        self.phonon.run_total_dos()

        _print("writing total dos")
        self.phonon.write_total_dos(filename=self.outputs.dos)

    def plot_bands(self) -> Tuple[Figure, Axes]:
        """run the default phonopy `plot_band_structure` method"""
        plt = self.phonon.plot_band_structure()
        return plt.gcf(), plt.gca()

A wrapper for running generic lammps scripts with injected variables Runs phonolammps and phonopy calculations on a given nanowire structure     and calculates the Γ->Χ phonon dispersion along the wire axis by default.     Adjust `phonon_path` and `phonon_mesh` accordingly to make changes.      The main output is an HDF5 file (default "band.hdf5").     With the default settings, this can easily be several GB in size for a     large (100-1000 atom) unit cell. # pylint: disable=too-many-instance-attributes # nothing works if these don't exist # validate paths before assigning \ PhononRunner(wire_datafile='{self.wire_datafile}',              potential_file='{self.potential_file}',              output_dir='{self.outputs.output_dir}',              generic_relax='{self.generic_relax}',              generic_in='{self.generic_in}',              marker='{self.outputs.marker}',              clean_up={self.clean_up},              trash_counter_max={self._trash_counter_max})                            self.phonon_path={self.phonon_path if self.phonon_path else None}              self.phonon_mesh={self.phonon_mesh if self.phonon_mesh else None}              self.nqpoints={self.nqpoints if self.nqpoints else None}              self.phl={self._phl if self._phl else None}              self.outputs={self.outputs} relax structure, compute forces, and compute phonons use LAMMPS to relax input structure         (minimize potential energy at 0 pressure) # aliasing for neatness # this file is written # IMPORTANT use phonolammps to compute and write the force constants # aliasing for neatness # re-write 'relaxed' data-file without doing minimization # TODO: fix these terrible names # delete substituted LAMMPS input file use phonolammps to compute and write the force constants # aliasing for neatness # re-write 'relaxed' data-file without doing minimization # delete substituted LAMMPS input file run the mesh on the phonon attribute use phonopy to compute band structure # use alternate function defined below use phonopy to compute the total DOS run the default phonopy `plot_band_structure` method