# -*- coding: utf-8 -*-

"""NVT (canonical) dynamics in LAMMPS"""

import lammps_step
import logging
import seamm
from seamm_util import ureg, Q_, units_class  # noqa: F401
import seamm_util.printing as printing
from seamm_util.printing import FormattedText as __
import random

logger = logging.getLogger(__name__)
job = printing.getPrinter()
printer = printing.getPrinter('lammps')

thermostat_metadata = {
    'Nose-Hoover':
        {
            'documentation': 'http://lammps.sandia.gov/doc/fix_nvt.html',
            'references': ['Shinoda', 'Tuckerman'],
        },
    'Berendsen':
        {
            'documentation':
                'http://lammps.sandia.gov/doc/fix_temp_berendsen.html',
            'references': ['Berendsen'],
        },
    'canonical sampling, velocity rescaling (csvr)':
        {
            'documentation': 'http://lammps.sandia.gov/doc/fix_temp_csvr.html',
            'references': ['Bussi1'],
        },
    'canonical sampling, langevin dynamics (csld)':
        {
            'documentation': 'http://lammps.sandia.gov/doc/fix_temp_csvr.html',
            'references': ['Bussi2'],
        },
    'velocity rescaling':
        {
            'documentation':
                'http://lammps.sandia.gov/doc/fix_temp_rescale.html',
            'references': [],
        },
    'Langevin':
        {
            'documentation': 'http://lammps.sandia.gov/doc/fix_langevin.html',
            'references': ['Schneider', 'Dunweg']
        },
}


class NVT(lammps_step.NVE):

    def __init__(self, flowchart=None, title='NVT dynamics', extension=None):
        """Initialize the node"""

        logger.debug('Creating NVT {}'.format(self))

        super().__init__(flowchart=flowchart, title=title, extension=extension)

        logger.debug('NVT after super init, {}'.format(self))

        self.description = 'NVT dynamics step in LAMMPS'

        logger.debug("NVT.init() creating NVT_Parameters object")

        self.parameters = lammps_step.NVT_Parameters()

        logger.debug("NVT.init() completed")

    def description_text(self, P=None):
        """Return a short description of this step.

        Return a nicely formatted string describing what this step will
        do.

        Keyword arguments:
            P: a dictionary of parameter values, which may be variables
                or final values. If None, then the parameters values will
                be used as is.
        """

        if not P:
            P = self.parameters.values_to_dict()

        # What will we do?

        if P['T0'] == P['T1']:
            text = "{time} of canonical (NVT) dynamics at {T0} "
        else:
            text = (
                "{time} of canonical (NVT) dynamics starting "
                " at {T0}, going to {T1}, "
            )
        if P['thermostat'] == 'Nose-Hoover':
            text += "using a Nose-Hoover thermostat."
            if P['Tchain'] != '3':
                if P['Tloop'] != '1':
                    text += (
                        " The thermostat will use a chain of {Tchain} "
                        "thermostats with {Tloop} subcycles and a "
                    )
                else:
                    text += (
                        " The thermostat will use a chain of {Tchain} "
                        "thermostats and a "
                    )
            elif P['Tloop'] != '1':
                text += " The thermostat will use {Tloop} subcycles and a "
            else:
                text += " The thermostat will use a "
            text += "drag factor of {drag}."
        elif P['thermostat'] == 'Berendsen':
            text += (
                "using a Berendsen thermostat with a damping time "
                "of {Tdamp}"
            )
        elif 'csvr' in P['thermostat']:
            text += (
                "using a canonical sampling thermostat using velocity "
                "rescaling (CSVR) with a damping time of {Tdamp} and "
                "a {seed}."
            )
        elif 'csld' in P['thermostat']:
            text += (
                "using a canonical sampling thermostat using Langevin "
                "dynamics (CSLD) with a damping time of {Tdamp} and "
                "a {seed}."
            )
        elif P['thermostat'] == 'velocity rescaling':
            text += (
                "using velocity rescaling every {frequency} with a "
                "temperature window of {window}."
            )
            if P['fraction'] != 1.0:
                text += (
                    " The velocities will only be scaled a fraction "
                    "({fraction}) of the amount needed to fully correct "
                    "the temperature."
                )
        elif P['thermostat'] == 'Langevin':
            text += (
                "using a Langevin thermostat with a damping time "
                "of {Tdamp} and a {seed}"
            )
        else:
            text += ("using the thermostat given by {thermostat}")

        return self.header + '\n' + __(text, **P, indent=4 * ' ').__str__()

    def describe(self, indent='', json_dict=None):
        """Write out information about what this node will do
        If json_dict is passed in, add information to that dictionary
        so that it can be written out by the controller as appropriate.
        """

        # Can't call super() because it will print too much
        self.visited = True
        job.job('\n' + self.indent + self.header)
        next_node = self.next()

        # Local copies of variables in a dictionary

        P = self.parameters.values_to_dict()
        text = self.description_text(P)
        job.job(__(text, indent=self.indent + '    ', **P))

        return next_node

    def get_input(self, extras=None):
        """Get the input for an NVT dynamics run in LAMMPS"""

        self.description = []

        P = self.parameters.current_values_to_dict(
            context=seamm.flowchart_variables._data
        )

        # Fix variables with special cases
        timestep, P['timestep'] = self.timestep(P['timestep'])

        if P['seed'] == 'random':
            # Apparently the seed must be no larger than 900,000,000
            P['seed'] = int(random.random() * 899999999)

        # Have to fix formatting for printing...
        PP = dict(P)
        for key in PP:
            if isinstance(PP[key], units_class):
                PP[key] = '{:~P}'.format(PP[key])

        self.description.append(
            __(self.description_text(PP), **PP, indent=3 * ' ').__str__()
        )

        time = lammps_step.to_lammps_units(P['time'], quantity='time')
        nsteps = round(time / timestep)

        T0 = lammps_step.to_lammps_units(P['T0'], quantity='temperature')
        T1 = lammps_step.to_lammps_units(P['T1'], quantity='temperature')
        Tdamp = lammps_step.to_lammps_units(P['Tdamp'], quantity='time')

        thermo_properties = (
            'time temp press etotal ke pe ebond '
            'eangle edihed eimp evdwl etail ecoul elong'
        )
        properties = 'v_time v_temp v_press v_etotal v_ke v_pe v_epair'
        title2 = 'tstep t T P Etot Eke Epe Epair'

        lines = []
        lines.append('')
        lines.append('#     NVT dynamics')
        lines.append('')
        lines.append('reset_timestep      0')
        lines.append('timestep            {}'.format(timestep))
        lines.append('thermo_style        custom {}'.format(thermo_properties))
        lines.append('thermo              {}'.format(int(nsteps / 100)))

        nfixes = 0
        if P['thermostat'] == 'Nose-Hoover':
            Tchain = P['Tchain']
            Tloop = P['Tloop']
            drag = P['drag']
            nfixes += 1
            lines.append(
                'fix                 {} all nvt '.format(nfixes) +
                'temp {} {} {} '.format(T0, T1, Tdamp) +
                'tchain {} '.format(Tchain) + 'tloop {} '.format(Tloop) +
                'drag {}'.format(drag)
            )
        elif P['thermostat'] == 'Berendsen':
            nfixes += 1
            lines.append(
                'fix                 {} '.format(nfixes) +
                'all temp/berendsen ' + ' {} {} {}'.format(T0, T1, Tdamp)
            )
            nfixes += 1
            lines.append('fix                 {} '.format(nfixes) + 'all nve')
        elif 'csvr' in P['thermostat']:
            seed = P['seed']
            nfixes += 1
            lines.append(
                'fix                 {} '.format(nfixes) + 'all temp/csvr ' +
                ' {} {} {} {}'.format(T0, T1, Tdamp, seed)
            )
            nfixes += 1
            lines.append('fix                 {} '.format(nfixes) + 'all nve')
        elif 'csld' in P['thermostat']:
            seed = P['seed']
            nfixes += 1
            lines.append(
                'fix                 {} '.format(nfixes) + 'all temp/csld ' +
                ' {} {} {} {}'.format(T0, T1, Tdamp, seed)
            )
            nfixes += 1
            lines.append('fix                 {} '.format(nfixes) + 'all nve')
        elif P['thermostat'] == 'velocity rescaling':
            frequency = lammps_step.to_lammps_units(
                P['frequency'], quantity='time'
            )

            nevery = round(nsteps / (frequency / timestep))
            window = lammps_step.to_lammps_units(
                P['window'], quantity='temperature'
            )
            fraction = P['fraction']
            nfixes += 1
            lines.append(
                'fix                 {} '.format(nfixes) +
                'all temp/rescale ' +
                '{} {} {} {} {}'.format(nevery, T0, T1, window, fraction)
            )
            nfixes += 1
            lines.append('fix                 {} '.format(nfixes) + 'all nve')
        elif P['thermostat'] == 'Langevin':
            seed = P['seed']
            nfixes += 1
            lines.append(
                'fix                 {} '.format(nfixes) + 'all langevin ' +
                '{} {} {} {} '.format(T0, T1, Tdamp, seed)
            )
            nfixes += 1
            lines.append('fix                 {} '.format(nfixes) + 'all nve')
        else:
            raise RuntimeError(
                "Don't recognize temperature control " +
                "'{}'".format(P['thermostat'])
            )

        # Add the citation for the thermostat
        for citation in thermostat_metadata[P['thermostat']]['references']:
            self.references.cite(
                raw=self._bibliography[citation],
                alias=citation,
                module='lammps_step',
                level=1,
                note='Citation for thermostat.'
            )

        # summary output written 10 times during run so we can see progress
        nevery = 10
        nfreq = int(nsteps / 10)
        nrepeat = int(nfreq / nevery)
        nfreq = nevery * nrepeat
        nfixes += 1
        lines.append(
            (
                "fix                 {} all ave/time {} {} {} {} off 2 "
                "title2 '{}' file summary_nvt_{}.txt"
            ).format(
                nfixes, nevery, nrepeat, nfreq, properties, title2,
                '_'.join(str(e) for e in self._id)
            )
        )
        # instantaneous output written for averaging
        if P['sampling'] == 'none':
            self.description.append(
                __(
                    "The run will be {nsteps:n} steps of dynamics.",
                    nsteps=nsteps,
                    indent=7 * ' '
                )
            )
        else:
            sampling = lammps_step.to_lammps_units(
                P['sampling'], quantity='time'
            )
            nevery = round(sampling / timestep)
            nfreq = int(nsteps / nevery)
            nrepeat = 1
            nfreq = nevery * nrepeat
            nfixes += 1

            if T0 == T1:
                title1 = (
                    '!MolSSI trajectory 1.0 LAMMPS, NVT {} steps of {} fs, '
                    'T={} K'
                ).format(int(nsteps / nevery), timestep * nevery, T0)
            else:
                title1 = (
                    '!MolSSI trajectory 1.0 LAMMPS, NVT {} steps of {} fs, '
                    'T={}-{} K'
                ).format(nsteps, timestep, T0, T1)
            lines.append(
                (
                    "fix                 {} all ave/time {} {} {} {} off 2"
                    " title1 '{}' title2 '{}' file trajectory_nvt_{}.seamm_trj"
                ).format(
                    nfixes, nevery, nrepeat, nfreq, properties, title1, title2,
                    '_'.join(str(e) for e in self._id)
                )
            )
            self.description.append(
                __(
                    (
                        "The run will be {nsteps:,d} steps of dynamics "
                        "sampled every {nevery:n} steps."
                    ),
                    nsteps=nsteps,
                    nevery=nevery,
                    indent=7 * ' '
                )
            )

        if extras is not None and 'shake' in extras:
            nfixes += 1
            lines.append(extras['shake'].format(nfixes))

        lines.append('')
        lines.append('run                 {}'.format(nsteps))
        lines.append('')

        for fix in range(1, nfixes + 1):
            lines.append('unfix               {}'.format(fix))
        lines.append('')

        return lines
