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

import math
import numpy as np
import operator

from .utils import Base

eps = 1e-4

class IsotopicFitRecord(object):
    """Describes a single isotopic pattern fit, comparing how well an
    experimentally observed sequence of peaks matches a theoretical isotopic
    pattern.

    IsotopicFitRecord instances are hashable and orderable (by score).

    Attributes
    ----------
    charge : int
        The charge state used to generate the theoretical pattern
    data : object
        An arbitrary Python object containing extra information
    experimental : list of FittedPeak
        The observed experimental peaks to be fitted
    missed_peaks : int
        The number of peaks in the theoretical pattern that do not
        have a matching experimental peak
    monoisotopic_peak : FittedPeak
        The fitted peak which corresponds to the monoisotopic peak
    score : float
        The score assigned to the fit by an IsotopicFitter object
    seed_peak : FittedPeak
        The peak that was used to initiate the fit. This may be unused
        if not using an Averagine method
    theoretical : list of TheoreticalPeak
        The theoretical isotopic pattern being fitted on the experimental data
    """
    __slots__ = ["seed_peak", "score", "charge", "experimental", "theoretical",
                 "monoisotopic_peak", "data", "missed_peaks"]

    def __init__(self, seed_peak, score, charge, theoretical, experimental, data=None,
                 missed_peaks=0, **kwargs):
        self.seed_peak = seed_peak
        self.score = score
        self.charge = charge
        self.experimental = experimental
        self.theoretical = theoretical
        self.monoisotopic_peak = experimental[0]
        self.data = data
        self.missed_peaks = missed_peaks

    def clone(self):
        return self.__class__(
            self.seed_peak, self.score, self.charge, self.theoretical, self.experimental, self.data, self.missed_peaks)

    def __reduce__(self):
        return self.__class__, (
            self.seed_peak, self.score, self.charge, self.theoretical, self.experimental, self.data, self.missed_peaks)

    def __eq__(self, other):
        val = (self.score == other.score and
               self.charge == other.charge and
               self.experimental == other.experimental and
               self.theoretical == other.theoretical)
        if self.data is not None or other.data is not None:
            val = val and (self.data == other.data)
        return val

    def __ne__(self, other):
        return not (self == other)

    def __lt__(self, other):
        return self.score < other.score

    def __gt__(self, other):
        return self.score > other.score

    def __hash__(self):
        return hash((self.monoisotopic_peak.mz, self.charge))

    def __iter__(self):
        yield self.score
        yield self.charge
        yield self.experimental
        yield self.theoretical

    @property
    def npeaks(self):
        return len(self.experimental)

    def __repr__(self):
        return "IsotopicFitRecord(score=%0.5f, charge=%d, npeaks=%d, monoisotopic_mz=%0.5f)" % (
            self.score, self.charge, self.npeaks, self.monoisotopic_peak.mz)

class FitSelectorBase(Base):
    """An object that controls the filtering and
    selection of IsotopicFitRecord

    Attributes
    ----------
    minimum_score : int
        The minimum score needed to be a candidate for selection. If the
        FitSelector is `minimizing` it is the maximal score to be a candidate.
    """
    minimum_score = 0

    def __init__(self, minimum_score=0):
        self.minimum_score = minimum_score

    def best(self, results):
        raise NotImplementedError()

    def __call__(self, *args, **kwargs):
        return self.best(*args, **kwargs)

    def reject(self, fit):
        raise NotImplementedError()

    def reject_score(self, score):
        raise NotImplementedError()

    def is_maximizing(self):
        return False

class MinimizeFitSelector(FitSelectorBase):
    """A FitSelector which tries to minimize the score of the best fit.
    """
    def best(self, results):
        """Returns the IsotopicFitRecord with the smallest score

        Parameters
        ----------
        results : list of IsotopicFitRecord
            List of isotopic fits to select the most optimal case from

        Returns
        -------
        IsotopicFitRecord
            The most optimal fit
        """
        return min(results, key=operator.attrgetter("score"))

    def reject(self, fit):
        """Decide whether the fit should be discarded for having too
        large a score. Compares against :attr:`minimum_score`

        Parameters
        ----------
        fit : IsotopicFitRecord

        Returns
        -------
        bool
        """
        return fit.score > self.minimum_score

    def reject_score(self, score):
        return score > self.minimum_score

    def is_maximizing(self):
        """Returns that this is *not* a maximizing selector

        Returns
        -------
        False
        """
        return False

class MaximizeFitSelector(FitSelectorBase):
    """A FitSelector which tries to maximize the score of the best fit.
    """
    def best(self, results):
        """Returns the IsotopicFitRecord with the largest score

        Parameters
        ----------
        results : list of IsotopicFitRecord
            List of isotopic fits to select the most optimal case from

        Returns
        -------
        IsotopicFitRecord
            The most optimal fit
        """
        return max(results, key=operator.attrgetter("score"))

    def reject(self, fit):
        """Decide whether the fit should be discarded for having too
        small a score. Compares against :attr:`minimum_score`

        Parameters
        ----------
        fit : IsotopicFitRecord

        Returns
        -------
        bool
        """
        return fit.score < self.minimum_score

    def reject_score(self, score):
        return score < self.minimum_score

    def is_maximizing(self):
        """Returns that this *is* a maximizing selector

        Returns
        -------
        False
        """
        return True

class IsotopicFitterBase(Base):
    '''A base class for Isotopic Pattern Fitters, objects
    which given a set of experimental peaks and a set of matching
    theoretical peaks, returns a fit score describing how good the
    match is.

    An IsotopicFitter may be optimal when the score is small (minimizing)
    or when the score is large (maximizing), and the appropriate
    :class:`FitSelectorBase` type will be used for :attr:`select`. This
    will also be reflected by :meth:`is_maximizing`.
    '''

    def __init__(self, score_threshold=0.5):
        self.select = MinimizeFitSelector(score_threshold)

    def evaluate(self, peaklist, observed, expected, **kwargs):
        """Evaluate a pair of peak lists for goodness-of-fit.

        Parameters
        ----------
        peaklist : :class:`~.PeakSet`
            The full set of all experimental peaks
        observed : list
            The list of experimental peaks that are part of this fit
        expected : list
            The list of theoretical peaks that are part of this fit
        **kwargs

        Returns
        -------
        float
            The score
        """
        return NotImplemented

    def _evaluate(self, peaklist, observed, expected, **kwargs):
        return self.evaluate(peaklist, observed, expected, **kwargs)

    def __call__(self, *args, **kwargs):
        """Invokes :meth:`evaluate`

        Parameters
        ----------
        *args
            Forwarded to :meth:`evaluate`
        **kwargs
            Forwarded to :meth:`evaluate`

        Returns
        -------
        float
            The score
        """
        return self.evaluate(*args, **kwargs)

    def reject(self, fit):
        """Test whether this fit is too poor to be used

        Parameters
        ----------
        fit : :class:`~IsotopicFitRecord`
            The fit to test

        Returns
        -------
        bool
        """
        return self.select.reject(fit)

    def is_maximizing(self):
        """Whether or not this fitter's score gets better as it grows

        Returns
        -------
        bool
            Whether or not this fitter is a maximizing fitter
        """
        return self.select.is_maximizing()

    def configure(self, deconvoluter, **kwargs):
        return self

class GTestFitter(IsotopicFitterBase):
    r"""Evaluate an isotopic fit using a G-test

    .. math::
        G = 2\sum_i^n{o_i * ({log}o_i - {log}e_i)}

    where :math:`o_i` is the intensity of the ith experimental peak
    and :math:`e_i` is the intensity of the ith theoretical peak.
    """
    def evaluate(self, peaklist, observed, expected, **kwargs):
        g_score = 2 * sum([obs.intensity * np.log(
            obs.intensity / theo.intensity) for obs, theo in zip(observed, expected)])
        return g_score

g_test = GTestFitter()

class ScaledGTestFitter(IsotopicFitterBase):
    r"""Evaluate an isotopic fit using a G-test after normalizing the
    list of experimental and theoretical peaks to both sum to 1.

    .. math::
        G = 2\sum_i^n{o_i * ({log}o_i - {log}e_i)}

    where :math:`o_i` is the intensity of the ith experimental peak
    and :math:`e_i` is the intensity of the ith theoretical peak.
    """
    def evaluate(self, peaklist, observed, expected, **kwargs):
        total_observed = sum(p.intensity for p in observed)
        total_expected = sum(p.intensity for p in expected)
        total_expected += eps
        normalized_observed = [obs.intensity /
                               total_observed for obs in observed]
        normalized_expected = [theo.intensity /
                               total_expected for theo in expected]
        g_score = 2 * sum([obs * np.log(obs / theo) for obs, theo in zip(
            normalized_observed, normalized_expected)])
        return g_score

g_test_scaled = ScaledGTestFitter()

class ChiSquareFitter(IsotopicFitterBase):

    def evaluate(self, peaklist, observed, expected, **kwargs):
        score = sum([(obs.intensity - theo.intensity)**2 / theo.intensity
                     for obs, theo in zip(observed, expected)])
        return score

chi_sqr_test = ChiSquareFitter()

class LeastSquaresFitter(IsotopicFitterBase):
    r"""Evaluate an isotopic fit using least squares coefficient of
    determination :math:`R^2`.

    .. math::
        {\hat e_i} &= e_i / max(e)

        {\hat t_i} &= t_i / max(t)

        {\hat t} &= \sum_i^n {\hat t_i}^2

        R^2 &= \frac{1}{{\hat t}}\sum_i^n ({\hat e_i} - {\hat t_i})^2

    where :math:`e_i` is the ith experimental peak intensity and :math:`t_i`
    is the ith theoretical peak intensity
    """

    def evaluate(self, peaklist, observed, expected, **kwargs):
        exp_max = max(p.intensity for p in observed)
        theo_max = max(p.intensity for p in expected)

        sum_of_squared_errors = 0
        sum_of_squared_theoreticals = 0

        for e, t in zip(observed, expected):
            normed_expr = e.intensity / exp_max
            normed_theo = t.intensity / theo_max
            sum_of_squared_errors += (normed_theo - normed_expr) ** 2
            sum_of_squared_theoreticals += normed_theo ** 2
        return sum_of_squared_errors / sum_of_squared_theoreticals

least_squares = LeastSquaresFitter()

class MSDeconVFitter(IsotopicFitterBase):
    '''An implementation of the scoring function used in :title-reference:`MSDeconV`

    References
    ----------
    Liu, X., Inbar, Y., Dorrestein, P. C., Wynne, C., Edwards, N., Souda, P., …
    Pevzner, P. A. (2010). Deconvolution and database search of complex tandem
    mass spectra of intact proteins: a combinatorial approach. Molecular & Cellular
    Proteomics : MCP, 9(12), 2772–2782. https://doi.org/10.1074/mcp.M110.002766
    '''

    def __init__(self, minimum_score=10, mass_error_tolerance=0.02):
        self.select = MaximizeFitSelector()
        self.select.minimum_score = minimum_score
        self.mass_error_tolerance = mass_error_tolerance

    def calculate_minimum_signal_to_noise(self, observed):
        snr = 0
        n = 0
        for obs in observed:
            if obs.signal_to_noise < 1:
                continue
            snr += obs.signal_to_noise
            n += 1
        return (snr / n) * 0.05

    def reweight(self, obs, theo, obs_total, theo_total):
        norm_obs = obs.intensity / obs_total
        norm_theo = theo.intensity / theo_total
        return norm_obs * np.log(norm_obs / norm_theo)

    def score_peak(self, obs, theo, mass_error_tolerance=0.02, minimum_signal_to_noise=1):
        if obs.signal_to_noise < minimum_signal_to_noise:
            return 0.

        mass_error = np.abs(obs.mz - theo.mz)

        if mass_error <= mass_error_tolerance:
            mass_accuracy = 1 - mass_error / mass_error_tolerance
        else:
            mass_accuracy = 0

        if obs.intensity < theo.intensity and (((theo.intensity - obs.intensity) / obs.intensity) <= 1):
            abundance_diff = 1 - \
                ((theo.intensity - obs.intensity) / obs.intensity)
        elif obs.intensity >= theo.intensity and (((obs.intensity - theo.intensity) / obs.intensity) <= 1):
            abundance_diff = np.sqrt(
                1 - ((obs.intensity - theo.intensity) / obs.intensity))
        else:
            abundance_diff = 0.
        score = np.sqrt(theo.intensity) * mass_accuracy * abundance_diff
        return score

    def evaluate(self, peaklist, observed, expected, **kwargs):
        score = 0
        for obs, theo in zip(observed, expected):
            inc = self.score_peak(obs, theo, self.mass_error_tolerance, 1)
            score += inc
        return score

class PenalizedMSDeconVFitter(IsotopicFitterBase):
    r'''An Isotopic Fitter which uses the :class:`MSDeconVFitter` score
    weighted by 1 - :attr:`penalty_factor` * :class:`ScaledGTestFitter` score

    .. math::
        S(e, t) = M(e, t) * (1 - G(e, t))

    where :math:`e` is the experimental peak list and :math:`t` is the theoretical
    peak list
    '''
    def __init__(self, minimum_score=10, penalty_factor=1., mass_error_tolerance=0.02):
        self.select = MaximizeFitSelector(minimum_score)
        self.msdeconv = MSDeconVFitter(mass_error_tolerance=mass_error_tolerance)
        self.penalizer = ScaledGTestFitter()
        self.penalty_factor = penalty_factor

    def evaluate(self, peaklist, observed, expected, **kwargs):
        score = self.msdeconv.evaluate(peaklist, observed, expected)
        penalty = abs(self.penalizer.evaluate(peaklist, observed, expected))
        return score * (1 - penalty * self.penalty_factor)

def decon2ls_chisqr_test(peaklist, observed, expected, **kwargs):
    fit_total = 0
    sum_total = 0
    for obs, theo in zip(observed, expected):
        intensity_diff = obs.intensity - theo.intensity
        fit_total += (intensity_diff ** 2) / (theo.intensity + obs.intensity)
        sum_total += theo.intensity * obs.intensity
    return fit_total / (sum_total + 0.01)

class InterferenceDetection(object):

    def __init__(self, peaklist):
        self.peaklist = peaklist

    def __call__(self, experimental_peaks, lower=None, upper=None):
        return self.detect_interference(experimental_peaks, lower, upper)

    def detect_interference(self, experimental_peaks, lower=None, upper=None):
        min_peak = experimental_peaks[0]
        max_peak = experimental_peaks[-1]

        if lower is None:
            try:
                lower = min_peak.mz - min_peak.full_width_at_half_max
            except AttributeError:
                lower = min_peak.mz
        if upper is None:
            try:
                upper = max_peak.mz + max_peak.full_width_at_half_max
            except AttributeError:
                upper = max_peak.mz

        region = self.peaklist.between(lower, upper)

        included_intensity = sum(p.intensity for p in experimental_peaks)
        region_intensity = sum(p.intensity for p in region)
        if region_intensity == 0:
            return 0.

        score = 1 - (included_intensity / region_intensity)
        return score

class DistinctPatternFitter(IsotopicFitterBase):

    def __init__(self, minimum_score=0.3, peak_count_scale=1.5, domain_scale=100.):
        self.select = MinimizeFitSelector(minimum_score)
        self.interference_detector = None
        self.g_test_scaled = ScaledGTestFitter()
        self.peak_count_scale = peak_count_scale
        self.domain_scale = domain_scale

    def evaluate(self, peaklist, experimental, theoretical, **kwargs):
        npeaks = float(len(experimental))
        if self.interference_detector is None:
            self.interference_detector = InterferenceDetection(peaklist)

        score = self.g_test_scaled(peaklist, experimental, theoretical)
        score *= abs((self.interference_detector.detect_interference(experimental) + 0.00001) / (
            npeaks * self.peak_count_scale)) * self.domain_scale
        return score

def percentile(N, percent):
    if not N:
        return None
    k = (len(N) - 1) * percent
    f = math.floor(k)
    c = math.ceil(k)
    if f == c:
        return N[int(k)]
    d0 = N[int(f)] * (c - k)
    d1 = N[int(c)] * (k - f)
    return d0 + d1

class DotProductFitter(IsotopicFitterBase):
    def evaluate(self, peaklist, observed, expected, **kwargs):
        total = 0
        for e, t in zip(observed, expected):
            total += e.intensity * t.intensity
        return total

try:
    _c = True
    _IsotopicFitRecord = IsotopicFitRecord
    _LeastSquaresFitter = LeastSquaresFitter
    _MSDeconVFitter = MSDeconVFitter
    _ScaledGTestFitter = ScaledGTestFitter
    _PenalizedMSDeconVFitter = PenalizedMSDeconVFitter
    _DistinctPatternFitter = DistinctPatternFitter
    _DotProductFitter = DotProductFitter

    from ._c.scoring import (
        IsotopicFitRecord, LeastSquaresFitter, MSDeconVFitter,
        ScaledGTestFitter, PenalizedMSDeconVFitter, DistinctPatternFitter,
        DotProductFitter)
except ImportError as e:
    print(e)
    _c = False

msdeconv = MSDeconVFitter()
least_squares = LeastSquaresFitter()
g_test_scaled = ScaledGTestFitter()
penalized_msdeconv = PenalizedMSDeconVFitter()
distinct_pattern_fitter = DistinctPatternFitter()

class MassShiftSupportPostProcessorBase(object):

    def __init__(self, shifts=None):
        if shifts is None:
            shifts = []
        self.shifts = shifts
        self._fits = None

    @property
    def fits(self):
        return self._fits

    @fits.setter
    def fits(self, value):
        self._fits = value

    def reweight(self, fit, shift):
        pass


#!/usr/bin/env python
# coding: utf-8

# In[1]:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import panel as pn

pn.extension("katex", "mathjax")

# ## Lecture 3-  Groundwater as a reservoir ##
#
# _(The contents presented in this section were re-developed principally by [Prof. Peter Dietrich](https://www.ufz.de/index.php?de=37303) and Dr. P. K. Yadav. The original contents are from Prof. Rudolf Liedl)_
#
# ---
#
# The content of the previous section was dedicated to very fundamental properties, such as aquifer and its types, solid and liquid (water) volumes in an aquifer, a of subsurface.
#
# In this lecture, the subsurface will be considered from the perspective as a groundwater reservoir and some key definition and parameters will be introduced.
#
#
# ### Groundwater and Aquifers
#
#
# If the subterranean water completely fills the pore space we call it groundwater. In Germany a slightly different definition is in use. There, groundwater is only the subterranean water which is not subject to other forces than gravity (see Fig. below). That means, that the water adhesively bound to the grains is not part of the groundwater. Applicable forces in groundwater are sometime locally defined. For, e.g., In **Germany** the _gravity_ is the only force acting on groundwater, whereas forces such adhesion, cohesion along with gravity are considered internationally.
#
#
#
# The difference between the international and the German definition of groundwater is the consideration of the adhesive water. Adhesive water does not participate in water movement. The same is true for water in isolated pores or in dead-end pores. All subterranean water not participating in water movement is summarized as immobile water. In contrast, the mobile water is the subterranean water participating in water movement.

# In[2]:

p1 = pn.pane.Markdown("""
Above the domain where the pore space is completely filled with water, there can may be an area with _partial water saturation._
This domain is termed as **unsaturated** or **vadose** or **aeration** zone. In this zone, the subterranean water is named vadose
or suspended water. The water content in the _unsaturated_ zone is quantified with the parameter **water saturation** **_(S)_** and can
be calculated as the ratio between the water-filled pore volume **_Vp,w_** and the total pore volume **_Vp_**, i.e.,
""",width=400, style={"text-align": "justify"})

p2 = pn.pane.LaTeX(r"""
$$
S = \frac{V_{p,w}}{V_p}
$$
""" )

p3 = pn.pane.Markdown("""
**_S = 1_** in the saturated zone. The boundary between the unsaturated and the saturated zone is named as the **groundwater table** or **groundwater surface**
(see Fig. at the side).
""")

video1 = pn.pane.Video("images/L03_f_4.mp4", width=600, height=400, loop=False)

p4 = pn.Column(p1, p2, p3)

pn.Row(p4, pn.Spacer(width=50), video1)

#

# The volumetric share of pore, which can be occupied by mobile water, is termed effective porosity $n_e$ or flow-through porosity
#
# $$
# n_e = \frac{V_{p, m}}{V_T}
# $$
#
# The effective porosity is dimensionless and the pore volume $V_{p,m}$ available for mobile water as well as the total volume $V_T$ has the dimension $L^3$. Effective porosity cannot exceed total porosity, i.e. $n_e \leq n$. The difference $(n – n_e)$ is termed **specific retention** or **field capacity**. Specific retention is the volumetric share of water which is retained in the porous medium after drainage due to gravitation.  The reason for retention is the adhesive force which bounds water at the grain surfaces. Because the available grain surface in a medium depends on the grain size, the effective porosity is also different for various materials. As shown in Figure (below) clay has a high total porosity but only a low effective porosity whereas the total porosity of cobbles is not so significant different from the effective porosity of this material.
#
#
#

# #### Example Problem
#
# Moist sand specimen = 72.5 cm$^3$ and its weight = 152 g
# Oven dried sample = 71.2 cm$^3$ and its weight = 145 g
#
# Other available information:
# Specific weight of particles $(\gamma_s)$ = 2.65 g/cm$^3$
# Specific weight of water $(\gamma_w)$ = 1 g/cm$^3$
#
# Find, total porosity, void ratio, water content, degree of saturation and effective porosity.

# In[3]:

# solution

V_ms = 72.5 # cm^3, volume moist sand
W_ms = 152 # g, weight moist sand, also total volume
V_ds = 71.2 # cm^3, volume dry sand
W_ds = 145 # g, weight dey sand

# Other info
g_s = 2.65 # g/cm^3, sp. weight, particles
g_w = 1 # g/cm^3, sp. wt. water

# intermediate calculation
V_w = (W_ms-W_ds)/g_w # cm^3,  W_w/g_w; density = M/V
V_s = W_ds/g_s # cm^3,
V_v = V_ms - V_s # cm^2, volume of voids
W_w = W_ms - W_ds # g, weight of water

# results calculation

n = V_v/V_ms*100 # %, Total porosity
e = V_v/V_s *100 # %,  void ratio
w = W_w/W_ds *100 # %, moisture content
S = V_w/V_v * 100 # %, degree of saturation

print("Total porosity is {0:0.2f}%".format(n))
print("Void ratio is {0:0.2f}%".format(e))
print("Moisture content is {0:0.2f}%".format(w))
print("Degree of saturation is {0:0.2f}%".format(S))

# In[4]:

p5 = pn.pane.Markdown(""" ###Aquifer Classifications """)

p6 = pn.pane.Markdown("""
Subterranean formations can be classified by the capability to store and/or transmit groundwater under natural conditions.

> An **aquifer** or a **groundwater reservoir** can store and transmit significant (= exploitable) amounts of groundwater.

> An **aquitard** can store and transmit groundwater but to a much lesser extent than an (adjacent) aquifer.

> An **aquiclude** can store groundwater but cannot transmit groundwater.

> An **aquifuge**can neither store nor transmit groundwater.

Aquifers usually appear as **layers**, i.e. their lateral extent is rather large as
compared to their vertical extent (maybe 10 – 100 km vs. 10 – 100 m). The upper aquifer
boundary is called **aquifer top**, and the lower boundary is called **aquifer bottom**.
The vertical distance between aquifer top and aquifer bottom is called **aquifer thickness**.
Upper and lower aquifer boundaries do not have to be horizontal and the thickness may be
spatially variable (see Fig. on the right).

 """, width=400, style={"text-align": "justify"})

p7 = pn.pane.Video("images/L03_f_6.mp4", width=600, height=200, loop=False)

p8 = pn.Column(p5, p6)
pn.Row(p8, pn.Spacer(width=25), p7)

# In[5]:

p9 = pn.pane.Markdown(""" ### Unconfined Aquifer """)

p10 = pn.pane.Markdown(""" An aquifer can be completely or only partially filled with groundwater (see Fig below).
Groundwater or an aquifer is termed **unconfined (phreatic)**, if the groundwater does not
extend up to the aquifer top. The position of the groundwater table is therefore changed
during water injection or extraction (“free“ groundwater table). Water in a borehole rises
up to the groundwater table.
""", width=600, style={"text-align": "justify"})

p11 = pn.pane.Video("images/L03_f_7.mp4", width=400, height=200, loop=False)
p12 = pn.pane.Markdown(""" ### Confined Aquifer """)

p13 = pn.pane.Markdown(""" Groundwater or an aquifer is termed **confined**, if the aquifer
contains groundwater throughout its entire thickness. The pore space remains completely
water filled during water injection or (moderate) extraction. This requires a low permeable
cover layer. In addition, the groundwater _recharge area_ must be located at higher altitude
than the aquifer top. The elevation of the _groundwater table_ in the recharge area defines
the position of the confined aquifer‘s pressure line. Water in a borehole rises up to the
pressure line (neglecting friction losses), i.e. higher than the elevation of the aquifer top.
""", width=600, style={"text-align": "justify"})

p14 = pn.pane.Video("images/L03_f_8.mp4", width=400, height=200, loop=False)

p15 = pn.pane.Markdown(""" ### Artesian Aquifer """)
p16 = pn.pane.Markdown(""" If the pressure line is above ground surface the we have an **artesian aquifer**
(see Fig. below). In this case, the water in a borehole would rises up to the ground
surface and then forms a fountain. _Artesian springs_ and _Artesian wells_ are based on this principle.
The type of an aquifer can change along a cross section.""", width=600, style={"text-align": "justify"})

p17 = pn.pane.Video("images/L03_f_9.mp4", width=400, height=200, loop=False)

pn.Column(p9, p10, p11, p12, p13, p14, p15, p16, p17)

# #### Example
#
# | Aquifer |  Obs. Point 1  | Obs. Point 2 | Obs. Point 3 | Obs. Point 4 |
# |:-------:|:--------------:|:------------:|:------------:|:------------:|
# |    A    |       ---      |      ---     |  Unconfined  |  Unconfined  |
# |    B    |       ---      |  Unconfined  |  Unconfined  |   Confined   |
# |    C    |   Unconfined   |  Unconfined  |   Confined   |   Confined   |
# |    D    | Conf./Artesian |   Confined   |   Confined   |   Confined   |
#
#
#
# _*perched aquifer_: Unconfined aquifer on top of another unconfined aquifer, separated from each other by a shallow aquitard

# ### Pressure and pressure head
# A reason for the movement of groundwater can be (hydrostatic) pressure difference. Let us consider a vertical column containing a porous medium and water filling the voids completely (Fig. below). We can assign to the top an arbitrary reference value pL for the pressure.  Due to the load of the overlaying water column, the pressure increase if we go deeper into the column. This is the same as what we observe if we are diving in a lake. The increase of the pressure depends on the density of the fluid and the depth below the water surface. In the setup given in Fig., we have the (hydrostatic) pressure $p$ [M/L/T$^2$] as a function of the height $z$ [L] above the reference point
#
# $$
# p(z) = p_L + \rho \cdot g \cdot (L - z)
# $$
#
# ```{margin} An optional title
# with, <br>
# $\rho$ [M/L$^2$] as the fluid density, <br>
# $g$ the acceleration of gravity [L/T$^2$], and  <br>
# L [L] the length of the water column
# ```
#
#
#
# As shown in Fig., we can add two observation points, one at the bottom ($z = 0$) and the other at the top of the column $(z=L)$. The pressure difference $\Delta p$ between the observation points is
#
# $$
# \Delta p = p(L) - p(0) = p_L - (p_L + \rho\cdot g \cdot L) = - \rho\cdot g \cdot L
# $$
#
# **Example:** Compare the pressure difference for an experimental setup (pipe length 50 cm) in which water and diesel are the two liquids.

# In[6]:

# solution
L_p = 50 # cm length of pipe
g = 981 # cm/s^2, accl. due to gravity

# Assume densities
rho_w = 1.0 # g/cm^3, density of water
rho_d = 0.830 # g/cm^3, density of diesel

# calculate
Dp_w = rho_w*g*L_p # g/cm.s^2, pressure difference due to water
Dp_d = rho_d*g*L_p# g/cm.s^2, pressure difference due to water
Dp_w
print("The pressure difference due to water is {0:2.2f} g/cm.s\u00b2,".format(Dp_w),
      "and that due to diesel is {0:0.2f} g/cm.s\u00b2.".format(Dp_d), )

# ### Hydrostatic pressure
#
# Mostly, pressure head is used instead of pressure when dealing with hydraulic properties or phenomena of the subsurface. The reason is that the pressure head can be easily measured with a tape whereas for a pressure measurement a more expensive manometer is necessary. The (hydrostatic) **pressure head** $\psi$ [L] is defined as
#
# $$
# \psi(z) = \psi_L + L - z
# $$
#
# This expression makes it clear why we can measure the pressure head with a tape. Similar to pressure head, the hydrostatic pressure can be schematically represented as
#
#

# ### Pressure in a Confined Aquifer
#
# Let us now consider a _confined aquifer_ (see Fig. below). The confining bed exerts a certain pressure $p_{cb}$ on the aquifer. This pressure is compensated partly by the porous medium and partly by the groundwater (pressures $p_{pm}$ and $p_w$, respectively). Therefore, we can write the following equation
#
# $$
# p_{cb} + p_{pm} + p_w = 0
# $$
#
#
#
# We can induce a _change in the hydrostatic pressure_ $\Delta p_w$ with an injection or release of groundwater. According to the above equation we can formulate
#
# $$
# \Delta p_{cb} + \Delta p_{pm} + \Delta p_w = 0.
# $$
#
# The hydrostatic pressure changes do not affect the weight of the confining bed and the exerted downward pressure remains unchanged. Therefore we have
#
# $$
# \Delta p_{cb} = 0 \:\:\: \text{and}
# $$
#
# $$
# \Delta p_{pm} = -\Delta p_w
# $$
#
# This implies that an increase of hydrostatic pressure automatically results in a decrease in the pressure exerted by the porous medium. In the case of a decrease of hydrostatic the pressure exerted by the porous medium would increase.
# The change in hydraulic pressure will have two effects with regard to water volume. First, the hydraulic pressure change $\Delta p_w$ directly leads to expansion/compression of the water and the water volume is accordingly increased/decreased. Secondly, the opposite change $\Delta p_{pm} = -\Delta p_w$ leads to compression/expansion of the porous medium as a whole (not the individual grains!). This, in turn, results in a reduced/an enlarged pore space such that the stored water volume is decreased/ increased. Both effects contribute to aquifer storage properties (see next section).

# ### Aquifer storage properties
#
# _Storage properties_ of the aquifer and associated parameters can be understood by considering pressure changes. For this purpose, we consider the effect of a _change in water volume_ $\Delta V_w'$ due to a _change in hydrostatic pressure_. The _relative_ changes in water volume $\Delta V_w'/\Delta w$ [-] are proportional to change of pressure in groundwater $\Delta p_w$:
#
# $$
# \frac{\Delta V_w'}{V_w} = \alpha_w \cdot \Delta p_w
# $$
#
# with $\alpha_w$ as the **compressibility of water** [LT$^2$/M]. The compressibility of water is roughly $4.4 \cdot 10^{-10}$ m$^2$/N. Taking into account an incompressible behavior of the water, that means an increase in hydrostatic pressure results in an inflow of water. A decrease in hydrostatic pressure results would cause an outflow of water. The above equation can be rearranged to yield
#
# $$
# \Delta V_w' = \alpha_w V_w \Delta p_w = \eta \alpha_w V_T \Delta p_w = \eta \alpha_w V_T \rho_w g \Delta \psi
# $$
# With $\eta$ [-] as the _total porosity_ and $\Delta \psi$ [L] as the _change in pressure head._
#

# ### Change in total volume
# The preceding considerations dealt with a change in storage by inflow or outflow. The change was invoked by a change of pressure in groundwater $\Delta p_w$. But a change could be also invoked by the change of the pressure exerted by the porous medium on the confining layer $\Delta p_{pm}$. A change $\Delta p_{pm}$ in the pressure results in a decrease or an increase $\Delta V_T$ in total aquifer volume. Both quantities are proportional to each other via
#
# $$
# \frac{\Delta V_T}{V_T} = - \alpha_{pm} \Delta p_{pm}
# $$
#
# whereby the ratio is the relative change of the total volume and $\alpha_{pm}$ the compressibility of the porous medium [LT$^2$/M]. The compressibility of the porous medium is roughly $10^{-10} - 10^{-8}$ m$^2$/N for gravel,
# $10^{-9} -  10^{-7}$ m$^2$/N for sand, and $10^{-8} - 10^{-6}$ m$^2$/N for clay. Taking into account the relation between pressure and pressure head, the above equation can be rearranged to yield
#
# $$
# \Delta V_T = -\alpha_{pm}V_T\Delta p_{pm} = \alpha_{pm}V_T\Delta p_{pw} = \alpha_{pm}V_T\rho_w g \Delta \psi
# $$
#
# $\Delta V_T$ represents a _change in volume of the porous medium_ as a whole. It is composed of a change in volume $\Delta V_s$ of the solids and another change $\Delta V_w''$ in water volume. Because the change in volume of the solid is negligible, we can write
#
# $$
# \Delta V_T = \Delta V_s + \Delta V_w'' \approx \Delta V_w''
# $$
#
# If we compare the last two equations we can immediately derive
#
# $$
# \Delta V_w'' = \alpha_{pm} V_T\rho_w g \Delta\psi
# $$
#
# With words that means a _decrease of pressure_ in the porous medium leads to _an expansion of the porous medium_ and an associated _increase in water volume_ and enlarged pore space. An _increase in pressure_ in the porous medium would lead to a _compression_ of the porous medium and an associated _decrease_ in water volume and _reduced_ pore space.
#
# The _total_ change $\Delta V_w$ in water volume consists of both effects caused by pressure changes $\Delta p_{pm}$ and $\Delta p_w$. Therefore we have
#
# $$
# \Delta V_w = \Delta V_w' + \Delta V_w''
# $$
#
# Using the results derived before, we can express how $\Delta V_w$ depends on changes $\Delta \psi$ in pressure head
#
# $$
# \Delta V_w = \Delta V_w' + \Delta V_w'' = \eta \alpha_w V_T \rho_w g \Delta \psi + \alpha_{pm} V_T \rho_w g \Delta \psi
# $$
#
# The _first term_ of the sum is related to changes in hydrostatic pressure $(\Delta p_w)$ and the _second term_ to pressure changes in the porous medium $(\Delta p_{pm})$.
#
# **Example:** The 45 m thick aquifer under the change of pressure 245 KPa compacts 0.20 m. What is the compressibility of porous media.

# In[7]:

# Available information

dP = 245 # KPa= KN/m^2, Change pressure
m = 45 # m, aquifer thickness
dm = 0.20 # m, change in aquifer thickness

print("V = A*m and dV = A * dm \n")

# Calculate
al_pm = (dm/m)/dP # m^2/KN, compressibility of porous media.

print("The compressibility of porous media in aquifer {0:0.2e} m\u00b2/KN.".format(al_pm))

# ### Specific storage
#
# For the characterization of the storage properties of an aquifer, we use the term **specific storage** $S_s$. It is defined as the volume of water that is released from a unit aquifer volume if hydrostatic pressure head is reduced by one unit
#
# $$
# S_s = \frac{\Delta V_w}{V_T \cdot \Delta \psi}
# $$
#
# The dimension of _specific storage_ is 1/L. Both impacts on water volume discussed before have to be considered in order to quantify $\Delta V_w$ in the above equation
#
# $$
# \Delta V_w = \eta \alpha_w V_T \rho_w g \Delta \psi + \alpha_{pm} V_T\rho_w g \Delta \psi
# $$
# The specific storage can therefore also be expressed as
#
# $$
# S_s = (\eta \alpha_w + \alpha_{pm})\rho_w g
# $$
#
# Typical values for specific storage range from $10^{-6}$ 1/m (e.g. gravel) to $10^{-2}$ 1/m (e.g. clay).
#

# ### Storativity
#
# Due to their relatively large lateral extent, aquifers are mostly considered as spatially two-dimensional (2D) systems. In this case, specific storage $S_s$ is replaced by the **storativity** or **storage coefficient** $S$ (Fig. below). Reference volume in a confined aquifer for defining specific storage $S_s$ is a unit cube (e.g. $V_T = 1$ m$^3$), and for defining storativity $S$ is a cuboid extending from the aquifer bottom to the aquifer top over a unit area (e.g. $A = 1$ m$^2$ and $V_T = A\cdot m$).
#
# For confined aquifers $S$ is simply obtained by multiplying $S_s$ by the aquifer thickness $m$
#
# $$
# S = S_s \cdot m
# $$
#
#
#
#
# _Storativity_ can be interpreted as the volume of water released from an aquifer volume extending from the aquifer bottom up to the aquifer top over a unit area if the hydrostatic pressure is reduced by one unit. _Storativity_ is dimensionless.
#
# Actually, **unconfined** aquifers are always treated as 2D systems. As a consequence, storativity is used to quantify water storage properties. The definition of storativity remains unchanged in principle but the considered aquifer volume now extends from the aquifer bottom up to the water table. For _unconfined aquifers_, storativity values correspond to _effective porosities_. This is explained by the free groundwater table. In this case a pressure changes simply lead to filling or emptying of voids. This is fundamentally different from the storage properties of confined aquifers discussed before.
#
# In **confined** aquifers all voids remain filled with groundwater during pressure changes and storage properties depend on the compressibilities of water and the porous medium.

# In[ ]:


from __future__ import division
from qm.dftbplus.dftbplus import DFTBplus
from qm.dftbplus.dftbpar import spin_w, spin_w_lc, max_l
from misc import au_to_A, call_name
import os, shutil, re, textwrap
import numpy as np

class SSR(DFTBplus):
    """ Class for SSR method of DFTB+

        :param object molecule: Molecule object
        :param boolean l_scc: Include self-consistent charge (SCC) scheme
        :param double scc_tol: Stopping criteria for the SCC iterations
        :param integer scc_max_iter: Maximum number of SCC iterations
        :param boolean l_onsite: Include onsite correction to SCC term
        :param boolean l_range_sep: Include long-range corrected functional
        :param string lc_method: Algorithms for LC-DFTB
        :param boolean l_ssr22: Use SSR(2,2) calculation?
        :param boolean l_ssr44: Use SSR(4,4) calculation?
        :param string guess: Initial guess method for SCC scheme
        :param string guess_file: Initial guess file for eigenvetors
        :param boolean l_state_interactions: Include state-interaction terms to SA-REKS
        :param double shift: Level shifting value in SCC iterations
        :param double,list tuning: Scaling factor for atomic spin constants
        :param string cpreks_grad_alg: Algorithms used in CP-REKS equations
        :param double cpreks_grad_tol: Tolerance used in the conjugate-gradient based algorithms
        :param boolean l_save_memory: Save memory in cache used in CP-REKS equations
        :param string embedding: Charge-charge embedding options
        :param boolean l_periodic: Use periodicity in the calculations
        :param double,list cell_length: The lattice vectors of periodic unit cell
        :param string sk_path: Path for Slater-Koster files
        :param string install_path: Path for DFTB+ install directory
        :param integer nthreads: Number of threads in the calculations
        :param string version: Version of DFTB+
    """
    def __init__(self, molecule, l_scc=True, scc_tol=1E-6, scc_max_iter=1000, l_onsite=False, \
        l_range_sep=False, lc_method="MatrixBased", l_ssr22=False, l_ssr44=False, guess="h0", \
        guess_file="./eigenvec.bin", l_state_interactions=False, shift=0.3, tuning=None, \
        cpreks_grad_alg="pcg", cpreks_grad_tol=1E-8, l_save_memory=False, embedding=None, \
        l_periodic=False, cell_length=[0., 0., 0., 0., 0., 0., 0., 0., 0.], sk_path="./", \
        install_path="./", nthreads=1, version="20.1"):
        # Initialize DFTB+ common variables
        super(SSR, self).__init__(molecule, sk_path, install_path, nthreads, version)

        # Initialize DFTB+ SSR variables
        self.l_scc = l_scc
        self.scc_tol = scc_tol
        self.scc_max_iter = scc_max_iter

        self.l_onsite = l_onsite
        if (self.l_onsite):
            error_message = "Onsite-correction not implemented!"
            error_vars = f"l_onsite = {self.l_onsite}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.l_range_sep = l_range_sep
        self.lc_method = lc_method.lower()

        self.l_ssr22 = l_ssr22
        # TODO : logical variable (l_ssr22, l_ssr44) must be changed for generality
        if (not self.l_ssr22):
            error_message = "Use (2,2) active space for SSR!"
            error_vars = f"l_ssr22 = {self.l_ssr22}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")
        self.l_ssr44 = l_ssr44
        if (self.l_ssr44):
            error_message = "(4,4) active space for SSR not implemented!"
            error_vars = f"l_ssr44 = {self.l_ssr44}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        # Set initial guess for eigenvectors
        self.guess = guess.lower()
        self.guess_file = guess_file
        if not (self.guess in ["h0", "read"]):
            error_message = "Invalid initial guess for DFTB/SSR!"
            error_vars = f"guess = {self.guess}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.l_state_interactions = l_state_interactions
        self.shift = shift

        # Set scaling factor for atomic spin constants
        self.tuning = tuning
        if (self.tuning != None):
            if (len(self.tuning) != len(self.atom_type)):
                error_message = "Number of elements for scaling factor must be equal to number of atom types!"
                error_vars = f"len(tuning) = {len(self.tuning)}, len(atom_type) = {len(self.atom_type)}"
                raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.cpreks_grad_alg = cpreks_grad_alg.lower()
        self.cpreks_grad_tol = cpreks_grad_tol
        self.l_save_memory = l_save_memory

        self.embedding = embedding
        if (self.embedding != None):
            self.embedding = self.embedding.lower()

        if not (self.embedding in [None, "mechanical", "electrostatic"]):
            error_message = "Invalid charge embedding for QM/MM calculation!"
            error_vars = f"embedding = {self.embedding}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.l_periodic = l_periodic
        self.a_axis = np.copy(cell_length[0:3])
        self.b_axis = np.copy(cell_length[3:6])
        self.c_axis = np.copy(cell_length[6:9])

        # Set 'l_nacme' and 're_calc' with respect to the computational method
        # DFTB/SSR can produce NACs, so we do not need to get NACME from CIoverlap
        # DFTB/SSR can compute the gradient of several states simultaneously.
        molecule.l_nacme = False
        self.re_calc = False

    def get_data(self, molecule, base_dir, bo_list, dt, istep, calc_force_only):
        """ Extract energy, gradient and nonadiabatic couplings from SSR method

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        self.copy_files(istep)
        super().get_data(base_dir, calc_force_only)
        self.write_xyz(molecule)
        self.get_input(molecule, istep, bo_list)
        self.run_QM(base_dir, istep, bo_list)
        self.extract_QM(molecule, bo_list)
        self.move_dir(base_dir)

    def copy_files(self, istep):
        """ Copy necessary scratch files in previous step

            :param integer istep: Current MD step
        """
        # Copy required files to read initial guess
        if (self.guess == "read" and istep >= 0):
            # After T = 0.0 s
            shutil.copy(os.path.join(self.scr_qm_dir, "eigenvec.bin"), \
                os.path.join(self.scr_qm_dir, "../eigenvec.bin.pre"))

    def get_input(self, molecule, istep, bo_list):
        """ Generate DFTB+ input files: geometry.gen, dftb_in.hsd

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Make 'geometry.gen' file
        os.system("xyz2gen geometry.xyz")
        if (self.l_periodic):
            # Substitute C to S in first line
            file_be = open('geometry.gen', 'r')
            file_af = open('tmp.gen', 'w')
            first_row = True
            for row in file_be:
                if (first_row):
                    row = f'{molecule.nat_qm} S\n'
                    first_row = False
                file_af.write(row)
            # Add gamma-point and cell lattice information
            geom_periodic = textwrap.dedent(f"""\
            {0.0:15.8f} {0.0:15.8f} {0.0:15.8f}
            {self.a_axis[0]:15.8f} {self.a_axis[1]:15.8f} {self.a_axis[2]:15.8f}
            {self.b_axis[0]:15.8f} {self.b_axis[1]:15.8f} {self.b_axis[2]:15.8f}
            {self.c_axis[0]:15.8f} {self.c_axis[1]:15.8f} {self.c_axis[2]:15.8f}
            """)
            file_af.write(geom_periodic)
            file_be.close()
            file_af.close()
            os.rename('tmp.gen', 'geometry.gen')

        # Make 'point_charges.xyz' file used in electrostatic charge embedding of QM/MM
        if (self.embedding == "electrostatic"):
            # Make 'point_charges.xyz' file
            input_geom_pc = ""
            for iat in range(molecule.nat_qm, molecule.nat):
                input_geom_pc += "".join([f"{i:15.8f}" for i in molecule.pos[iat] * au_to_A])
                input_geom_pc += f"  {molecule.mm_charge[iat - molecule.nat_qm]:8.4f}\n"

            # Write 'point_charges.xyz' file
            file_name = "point_charges.xyz"
            with open(file_name, "w") as f:
                f.write(input_geom_pc)

        # Make 'dftb_in.hsd' file
        input_dftb = ""

        # Geometry Block
        input_geom = textwrap.dedent(f"""\
        Geometry = GenFormat{{
          <<< 'geometry.gen'
        }}
        """)
        input_dftb += input_geom

        # Hamiltonian Block
        input_ham_init = textwrap.dedent(f"""\
        Hamiltonian = DFTB{{
        """)
        input_dftb += input_ham_init

        # SCC-DFTB option
        if (self.l_scc):
            input_ham_scc = textwrap.indent(textwrap.dedent(f"""\
              SCC = Yes
              SCCTolerance = {self.scc_tol}
              MaxSCCIterations = {self.scc_max_iter}
            """), "  ")
            input_dftb += input_ham_scc

            # Read atomic spin constants used in DFTB/SSR
            if (self.l_range_sep):
                spin_constant = ("\n" + " " * 14).join([f"  {itype} = {{ {spin_w_lc[f'{itype}']} }}" for itype in self.atom_type])
            else:
                spin_constant = ("\n" + " " * 14).join([f"  {itype} = {{ {spin_w[f'{itype}']} }}" for itype in self.atom_type])
            input_ham_spin = textwrap.indent(textwrap.dedent(f"""\
              SpinConstants = {{
                ShellResolvedSpin = Yes
              {spin_constant}
              }}
            """), "  ")
            input_dftb += input_ham_spin

            # Long-range corrected DFTB (LC-DFTB) option
            if (self.l_range_sep):
                input_ham_lc = textwrap.indent(textwrap.dedent(f"""\
                  RangeSeparated = LC{{
                    Screening = {self.lc_method}{{}}
                  }}
                """), "  ")
                input_dftb += input_ham_lc

            # Add point charges to Hamiltonian used in electrostatic charge embedding of QM/MM
            if (self.embedding == "electrostatic"):
                input_ham_pc = textwrap.indent(textwrap.dedent(f"""\
                  ElectricField = PointCharges{{
                    CoordsAndCharges [Angstrom] = DirectRead{{
                      Records = {molecule.nat_mm}
                      File = "point_charges.xyz"
                    }}
                  }}
                """), "  ")
                input_dftb += input_ham_pc

        if (self.l_periodic):
            input_ham_periodic = textwrap.indent(textwrap.dedent(f"""\
              KPointsAndWeights = {{
                0.0 0.0 0.0 1.0
              }}
            """), "  ")
            input_dftb += input_ham_periodic

        angular_momentum = ("\n" + " " * 10).join([f"  {itype} = '{max_l[f'{itype}']}'" for itype in self.atom_type])
        input_ham_basic = textwrap.dedent(f"""\
          Charge = {molecule.charge}
          MaxAngularMomentum = {{
          {angular_momentum}
          }}
          SlaterKosterFiles = Type2FileNames{{
            Prefix = '{self.sk_path}'
            Separator = '-'
            Suffix = '.skf'
            LowerCaseTypeName = No
          }}
        }}
        """)
        input_dftb += input_ham_basic

        # Analysis Block
        input_analysis = textwrap.dedent(f"""\
        Analysis = {{
          CalculateForces = Yes
          WriteBandOut = Yes
          WriteEigenvectors = Yes
          MullikenAnalysis = Yes
        }}
        """)
        input_dftb += input_analysis

        # Options Block
        input_options = textwrap.dedent(f"""\
        Options = {{
          WriteDetailedXml = Yes
          WriteDetailedOut = Yes
          TimingVerbosity = -1
        }}
        """)
        input_dftb += input_options

        # REKS Block

        # Energy functional options
        if (self.l_ssr22):
            # Set active space and energy functional used in SA-REKS
            space = "SSR22"
            if (molecule.nst == 1):
                energy_functional = "{ 'PPS' }"
                all_states = "No"
            elif (molecule.nst == 2):
                energy_functional = "{ 'PPS' 'OSS' }"
                all_states = "No"
            elif (molecule.nst == 3):
                energy_functional = "{ 'PPS' 'OSS' }"
                all_states = "Yes"
            else:
                error_message = "SSR(2,2) can calculate up to 3 electronic states!"
                error_vars = f"Molecule.nstates = {molecule.nst}"
                raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

            # Include state-interaction terms to SA-REKS; SI-SA-REKS
            if (self.l_state_interactions):
                do_ssr = "Yes"
            else:
                do_ssr = "No"

            # Read 'eigenvec.bin' from previous step
            if (self.guess == "read"):
                if (istep == -1):
                    if (os.path.isfile(self.guess_file)):
                        # Copy guess file to currect directory
                        shutil.copy(self.guess_file, os.path.join(self.scr_qm_dir, "eigenvec.bin"))
                        restart = "Yes"
                    else:
                        restart = "No"
                elif (istep >= 0):
                    # Move previous file to currect directory
                    os.rename("../eigenvec.bin.pre", "./eigenvec.bin")
                    restart = "Yes"
            elif (self.guess == "h0"):
                restart = "No"

            # Scale the atomic spin constants
            if (self.tuning != None):
                spin_tuning = ""
                spin_tuning += "{"
                for scale_W in self.tuning:
                    spin_tuning += f" {scale_W} "
                spin_tuning += "}"
            else:
                spin_tuning = "{}"

        # CP-REKS algorithm options
        if (self.cpreks_grad_alg == "pcg"):
            cpreks_alg = "ConjugateGradient"
            preconditioner = "Yes"
        elif (self.cpreks_grad_alg == "cg"):
            cpreks_alg = "ConjugateGradient"
            preconditioner = "No"
        elif (self.cpreks_grad_alg == "direct"):
            cpreks_alg = "Direct"
            preconditioner = "No"
        else:
            error_message = "Invalid algorithms for CP-REKS problem!"
            error_vars = f"cpreks_grad_alg = {self.cpreks_grad_alg}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        # Save memory in cache to reduce computational cost in CP-REKS
        if (self.l_save_memory):
            memory = "Yes"
        else:
            memory = "No"

        # NAC calculation options
        if (molecule.nst == 1 or not self.l_state_interactions):
            # Single-state REKS or SA-REKS state
            self.nac = "No"
        else:
            # SSR state
            if (self.calc_coupling):
                # SHXF, SH, Eh need NAC calculations
                self.nac = "Yes"
            else:
                # BOMD do not need NAC calculations
                self.nac = "No"

        # Relaxed density options; It is determined automatically
        if (molecule.l_qmmm and self.embedding == "electrostatic"):
            relaxed_density = "Yes"
        else:
            relaxed_density = "No"

        input_reks = textwrap.dedent(f"""\
        REKS = {space}{{
          Energy = {{
            Functional = {energy_functional}
            StateInteractions = {do_ssr}
            IncludeAllStates = {all_states}
          }}
          TargetState = {bo_list[0] + 1}
          ReadEigenvectors = {restart}
          FONmaxIter = 50
          shift = {self.shift}
          SpinTuning = {spin_tuning}
          Gradient = {cpreks_alg}{{
            CGmaxIter = 100
            Tolerance = {self.cpreks_grad_tol}
            Preconditioner = {preconditioner}
            SaveMemory = {memory}
          }}
          RelaxedDensity = {relaxed_density}
          NonAdiabaticCoupling = {self.nac}
          VerbosityLevel = 1
        }}
        """)
        input_dftb += input_reks

        # ParserOptions Block
        if (self.version == "19.1"):
            error_message = "SSR not implemented in this version!"
            error_vars = f"version = {self.version}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")
        elif (self.version == "20.1"):
            parser_version = 8

        input_parseroptions = textwrap.dedent(f"""\
        ParserOptions = {{
          ParserVersion = {parser_version}
        }}
        """)
        input_dftb += input_parseroptions

        # Write 'dftb_in.hsd' file
        file_name = "dftb_in.hsd"
        with open(file_name, "w") as f:
            f.write(input_dftb)

    def run_QM(self, base_dir, istep, bo_list):
        """ Run SSR calculation and save the output files to qm_log directory

            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Set run command
        qm_command = os.path.join(self.qm_path, "dftb+")
        # OpenMP setting
        os.environ["OMP_NUM_THREADS"] = f"{self.nthreads}"
        command = f"{qm_command} > log"
        # Run DFTB+ method for molecular dynamics
        os.system(command)

        # Copy the output file to 'qm_log' directory
        tmp_dir = os.path.join(base_dir, "qm_log")
        if (os.path.exists(tmp_dir)):
            log_step = f"log.{istep + 1}.{bo_list[0]}"
            shutil.copy("log", os.path.join(tmp_dir, log_step))

    def extract_QM(self, molecule, bo_list):
        """ Read the output files to get BO information

            :param object molecule: Molecule object
            :param integer,list bo_list: List of BO states for BO calculation
        """
        # Read 'log' file
        file_name = "log"
        with open(file_name, "r") as f:
            log_out = f.read()

        # Read 'detailed.out' file
        file_name = "detailed.out"
        with open(file_name, "r") as f:
            detailed_out = f.read()

        # Energy
        if (molecule.nst == 1):
            # Single-state REKS
            tmp_e = 'Spin' + '\n\s+\w+\s+([-]\S+)(?:\s+\S+){3}' * molecule.nst
            energy = re.findall(tmp_e, log_out)
            energy = np.array(energy, dtype=np.float64)
        else:
            if (self.l_state_interactions):
                # SSR state
                energy = re.findall('SSR state\s+\S+\s+([-]\S+)', log_out)
                energy = np.array(energy, dtype=np.float64)
            else:
                # SA-REKS state
                tmp_e = 'Spin' + '\n\s+\w+\s+([-]\S+)(?:\s+\S+){3}' * molecule.nst
                energy = re.findall(tmp_e, log_out)
                energy = np.array(energy[0], dtype=np.float64)

        for ist in range(molecule.nst):
            molecule.states[ist].energy = energy[ist]

        # Force
        if (self.nac == "Yes"):
            # SHXF, SH, Eh : SSR state
            if (molecule.l_qmmm):
                for ist in range(molecule.nst):
                    tmp_g = f' {ist + 1} st state \(SSR\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                    grad = re.findall(tmp_g, log_out)
                    grad = np.array(grad[0], dtype=np.float64)
                    grad = grad.reshape(molecule.nat_qm, 3, order='C')
                    molecule.states[ist].force[0:molecule.nat_qm] = - grad
                    if (self.embedding == "electrostatic"):
                        tmp_f = 'Forces on external charges' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_mm
                        force = re.findall(tmp_f, detailed_out)
                        force = np.array(force[0], dtype=np.float64)
                        force = force.reshape(molecule.nat_mm, 3, order='C')
                        molecule.states[ist].force[molecule.nat_qm:molecule.nat] = force
            else:
                for ist in range(molecule.nst):
                    tmp_g = f' {ist + 1} st state \(SSR\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                    grad = re.findall(tmp_g, log_out)
                    grad = np.array(grad[0], dtype=np.float64)
                    grad = grad.reshape(molecule.nat_qm, 3, order='C')
                    molecule.states[ist].force = - np.copy(grad)
        else:
            # BOMD : SSR state, SA-REKS state or single-state REKS
            if (molecule.l_qmmm):
                tmp_g = f' {bo_list[0] + 1} state \(\w+[-]*\w+\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                grad = re.findall(tmp_g, log_out)
                grad = np.array(grad[0], dtype=np.float64)
                grad = grad.reshape(molecule.nat_qm, 3, order='C')
                molecule.states[bo_list[0]].force[0:molecule.nat_qm] = - grad
                if (self.embedding == "electrostatic"):
                    tmp_f = 'Forces on external charges' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_mm
                    force = re.findall(tmp_f, detailed_out)
                    force = np.array(force[0], dtype=np.float64)
                    force = force.reshape(molecule.nat_mm, 3, order='C')
                    molecule.states[bo_list[0]].force[molecule.nat_qm:molecule.nat] = force
            else:
                tmp_g = f' {bo_list[0] + 1} state \(\w+[-]*\w+\)' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                grad = re.findall(tmp_g, log_out)
                grad = np.array(grad[0], dtype=np.float64)
                grad = grad.reshape(molecule.nat_qm, 3, order='C')
                molecule.states[bo_list[0]].force = - np.copy(grad)

        # NAC
        if (self.nac == "Yes"):
            kst = 0
            for ist in range(molecule.nst):
                for jst in range(ist + 1, molecule.nst):
                    tmp_c = 'non-adiabatic coupling' + '\n\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
                    nac = re.findall(tmp_c, log_out)
                    nac = np.array(nac[kst], dtype=np.float64)
                    nac = nac.reshape(molecule.nat_qm, 3, order='C')
                    molecule.nac[ist, jst] = nac
                    molecule.nac[jst, ist] = - nac
                    kst += 1


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# pylama:ignore=D213,D400,E501,E221,D103,D205,D209,E124,E127,E128

"""Find optimal nutritional combinations

This optimizer assumes that the foods to be targeted are in the file
foods.jsontxt.gz as created by build_foods.sh. That means that each line is the
JSON representation of a record as created by nndb-import (see
https://github.com/CraigKelly/nndb-import)
"""

import sys
import gzip

from functools import partial

# Try to use various speedy json libraries then fall back to stdlib
try:
    import ujson as json
except:
    try:
        sys.stderr.write("No ujson, trying cjson\n")
        import cjson as json
    except:
        sys.stderr.write("No cjson, using stdlib json\n")
        import json

import scipy
import scipy.stats
sp = scipy
stats = sp.stats

from micro_reqs import ALL_MICROS  #NOQA

# Parameters we use
INIT_RANDOMS    = 10000                       # Initial random members of population
RND_EXP_ENTRIES = 4                           # Exp value for # of non-zero entries
MAX_AMT         = 5.0                         # Non-zero entry max value
POP_SIZE        = 200                         # Pop size (after initial population)
MUTATE_RATE     = 0.20                        # Rate at which entry is mutated
MUTATE_BIG      = MUTATE_RATE * 0.20          # Part of mutate rate, chance of "big" mutation
IMMORTALS       = 5                           # Number of pop members that carry over
MERGE_SIZE      = 4                           # Number of solutions we sample for merge
SPARSITY_MAX    = (RND_EXP_ENTRIES-1) * (MAX_AMT*0.75)

# We pre-define random routines so that we can swap out (Python's native
# Mersenne twister is fine but a little slow for what we're doing)
randrange = partial(sp.random.randint, 0)
rand = sp.random.rand

def choice(lst):
    return lst[randrange(len(lst))]

# Simple logistic function for scipy arrays
def logistic(spa):
    return 1.0 / (1.0 + scipy.exp(-spa))

# Simple logging stand-in
def info(s, *args):
    if args:
        s = s % args
    print(s)

FILTER_WORDS = [i.replace('\s', ' ') for i in """
    applesauce        mashed
    bacon             meatless
    beverage          millet
    breaded           nectar
    bulgur            noodles
    candied           oil-roasted
    catsup            pork
    celery\sflakes
    chili             products
    chowchow          puffs
    cocktail          ranch
    cilantro
    cornmeal          relish
    cornstarch        rice
    couscous          soymilk
    dehydrated        spaghetti
    flour             spread
    franks            succotash
    fried             syrup
    groats            tapioca
    hash brown        taro
    hominy            tomato\sproducts
    hummus            vegetables
    juice             vegetarian
    lambsquarters     veggie
    liquid\sfrom      vermicelli
    macaroni          wheat
""".split()]

FILTER_SPLITS = ["sweet", "sweetened" "wheat", ]

def food_filter(line):
    line = line.strip()
    if not line:
        return None

    food = json.loads(line)
    if not food:
        return None

    descrip = food.get('descrip', '').lower()
    if not descrip:
        return None

    for word in FILTER_WORDS:
        if word in descrip:
            return None

    parts = descrip.split(',')
    for part in parts:
        part = part.strip()
        if not part:
            continue
        for check in FILTER_SPLITS:
            if check == part:
                return None

    return food

def read_foods(filename):
    # Test show zcat subprocess MUCH faster then `with gzip.open(filename, 'r')`
    with gzip.open(filename) as fh:
        for line in fh:
            food = food_filter(line.decode('utf-8'))
            if food:
                yield food

def extract_nutrients(food):
    contains = dict(
        (n["nutrient_id"], n["nutrient_val"]) for n in food["nutrients"]
    )
    return [contains.get(nid, 0.0) for _, nid, _, _, _ in ALL_MICROS]

class opt_engine(object):
    """Given a list of food objects from our file, perform the various
    optmization functions we want
    """

    def __init__(self, food_details):
        """Load initial matrices and create initial expanded population."""
        self.food_details = food_details
        self.food_count = len(food_details)

        # We want our random instances to have an expected number of entries
        # equal to RND_EXP_ENTRIES, which for Binomial(n,p) is np. So we know
        # p is RND_EXP_ENTRIES/n
        self.rnd_entries = partial(
            stats.binom.rvs,
            self.food_count,
            float(RND_EXP_ENTRIES) / float(self.food_count)
        )

        # |foods| x |nutrients| 2-d array
        self.foods = sp.array([extract_nutrients(f) for f in food_details])

        # RDA and UL vectors
        self.RDA, self.UL = [], []
        for _, _, _, r, u in ALL_MICROS:
            self.RDA.append(r)
            self.UL.append(u)
        self.RDA, self.UL = sp.array(self.RDA), sp.array(self.UL)

        # Our first generation contains all single foods and INIT_RANDOMS
        # inital random entries
        pop = []
        for i in range(self.food_count):
            one = [0.0] * self.food_count
            one[i] = 1.0
            pop.append(list(one))

        for _ in range(INIT_RANDOMS):
            pop.append(self.rand_inst())

        self.set_pop(pop)
        self.last_population = []
        self.last_popnutrition = []

    def set_pop(self, pop):
        """Directly set population and recalc the pop-nutrition grid."""
        # |members| x |foods| 2-d array
        self.population = sp.array(pop)
        self.popnutrition = sp.dot(self.population, self.foods)

    def rand_inst(self):
        """Return a random instance based on RND_EXP_ENTRIES and current foods."""
        inst = [0.0] * self.food_count

        entries = self.rnd_entries()
        if entries < 1:
            entries = RND_EXP_ENTRIES  # extra boost to our average

        idxs = set([randrange(len(inst))])
        while len(idxs) < entries:
            idxs.add(randrange(len(inst)))

        for idx in idxs:
            inst[idx] = choice([0.5, 1.0, 1.5, 2.0])

        return inst

    def calories(self, inst):
        """Given an instance, return the total calories and the total g's of sugar."""
        total_cals = 0.0
        total_sugar = 0.0

        for idx, amt in enumerate(inst):
            if abs(amt - 0.0) < 0.00001:
                continue

            food = self.food_details[idx]

            for n in food["nutrients"]:
                if n["nutrient_id"] == "208":
                    total_cals += float(n["nutrient_val"]) * amt
                if n["nutrient_id"] == "269":
                    total_sugar += float(n["nutrient_val"]) * amt

        return total_cals, total_sugar

    SCORE_WEIGHTS = sp.array([
        2.0,  # RDA
        1.0,  # nutrient per gram
        1.8,  # UL
        1.5,  # cals/gram
        1.2,  # sparsity
        1.0,  # sugar
    ])

    def score(self, inst, nutr):
        """Return a fitness score and the originating score vector for the given
        instance and it's nutrition space vector."""
        assert len(inst) == self.foods.shape[0]
        assert len(nutr) == self.foods.shape[1]
        assert len(self.RDA) == len(self.UL) == len(nutr)

        # Figure any nutrients above the upper limit in percentage units.
        too_much = (nutr - self.UL) / self.UL
        calories, sugar = self.calories(inst)
        grams = 100.0 * inst.sum()  # Our database measures everything in 100g units

        # nutrient as percent of RDA, clipped to remove advantage for going over
        nutr_scores = sp.clip(nutr / self.RDA, 0.0, 1.01)

        components = self.SCORE_WEIGHTS * sp.array([
            # good stuff (note that we want to MINIMIZE our score
            0.0 - nutr_scores.mean(),                                # RDA, but no credit for going over
            0.0 - (nutr_scores.sum() / grams),                       # Nutrient score per gram

            # bad stuff
            too_much[too_much > -0.01].sum(),                        # UL, but only as we get high
            logistic(calories / grams),                              # cals per gram
            logistic(len(inst[inst > 0]) / float(RND_EXP_ENTRIES)),  # Simpler is better
            logistic(sugar)                                          # sugar is bad, mmm-kay
        ])

        return components.mean(), components

    def mutate(self, inst):  #NOQA
        """Return a mutated copy of the given instance"""
        rnd_inst = list(inst)

        for idx, val in enumerate(inst):
            roll = rand()
            if roll > MUTATE_RATE:
                continue  # No mutation here

            if roll < MUTATE_BIG and val > 0.0:
                # Swapping with blank index - find the other, do the swap, and leave
                blank_idx = randrange(len(inst))
                while inst[blank_idx] > 0.0:
                    blank_idx = randrange(len(inst))
                rnd_inst[idx], rnd_inst[blank_idx] = rnd_inst[blank_idx], rnd_inst[idx]
                continue

            if roll < MUTATE_BIG:
                modifier = choice([1.5, 1.75, 2.0, 2.25])  # Big jump
            else:
                modifier = choice([0.25, 0.5, .75, 1.0, 1.25])  # Regular jump

            if val <= 0.0:
                val = modifier
            else:
                val = sp.clip(val + (choice([-1.0, 1.0]) * modifier), 0.0, MAX_AMT)

            rnd_inst[idx] = val

        # All of this is great, but we tend to just get bigger and bigger - we want
        # lots of sparsity AND a low total sum
        while sum(rnd_inst) > SPARSITY_MAX:
            for idx in range(len(rnd_inst)):
                if rnd_inst[idx] <= 0.1:
                    rnd_inst[idx] = 0.0
                elif rand() < 0.70:
                    rnd_inst[idx] *= 0.5
                else:
                    rnd_inst[idx] = 0.0

        return rnd_inst

    def crossover(self, inst1, inst2):
        """Return a new instance via cross over with the 2 given instances."""
        return [x if rand() < 0.5 else y for x, y in zip(inst1, inst2)]

    def winner(self):
        """Currently, chose a winner from the population with tournament sel."""
        idx1 = randrange(len(self.population))
        idx2 = randrange(len(self.population))
        return min(idx1, idx2)

    def generation(self):
        """Do a single generation of work."""
        all_scored = []
        for inst, nutr in zip(self.population, self.popnutrition):
            score, score_vector = self.score(inst, nutr)
            all_scored.append((score, inst))

        all_scored.sort(key=lambda one: one[0])

        # Keep previous generation in sorted order
        self.last_population = sp.array([inst for _, inst in all_scored])
        self.last_popnutrition = sp.dot(self.last_population, self.foods)

        # Now we are creating the next generation

        # Don't add dups to a generation
        new_gen = []
        seen = set()

        def add_to_gen(inst):
            key = tuple(inst)
            if key in seen:
                return
            seen.add(key)
            new_gen.append(inst)

        # Keep top N entries, top N entries with mutations, and top N entries with
        # one-off removed and then mutated
        for score, inst in all_scored[:IMMORTALS]:
            # top entries
            add_to_gen(inst)
            # mutated
            add_to_gen(self.mutate(inst))
            # one-off removed and then mutated
            for idx, val in enumerate(inst):
                if val > 0.0:
                    one_off = list(inst)
                    one_off[idx] = 0.0
                    add_to_gen(one_off)
                    add_to_gen(self.mutate(one_off))

        # Fill out the rest of the generation
        # Note we're actually TRIPLE our population size - We use the population
        # size for new random entries, tradition select/crossover/mutate entries,
        # and sample/sum/mutate entries. (NOte our sample/sum/mutate works because
        # our mutate routine has a final "re-balance + force sparsity step"
        for _ in range(POP_SIZE):
            # simple random entry
            add_to_gen(self.rand_inst())

            # Merged entry
            merged = sp.array(self.population[self.winner()])
            for _ in range(MERGE_SIZE-1):
                merged += sp.array(self.population[self.winner()])
            add_to_gen(self.mutate(list(merged)))

            # "Normal" creation via selection, crossover/breeding. and mutation
            add_to_gen(self.mutate(self.crossover(
                self.population[self.winner()],
                self.population[self.winner()]
            )))

        # Have a new generation
        self.set_pop(new_gen)

    def dump_instance(self, inst, nutr, info=info):
        """Given an instance and it's nutritional space vector (just like score),
        output a decent, printable representation of the instance
        """
        score, score_vector = self.score(inst, nutr)
        info("Solution with score %12.4f (vector %s)", score, score_vector)
        for mic, amt in zip(ALL_MICROS, nutr):
            name, _, units, rda, ul = mic
            info("  NUTR %-15s: %8.2f %-5s, %8.2f%% of RDA %8.2f (UL %8.2f)",
                         name,  amt, units, (amt/rda)*100.0, rda,      ul
            )

        total_cals = 0.0

        for idx, amt in enumerate(inst):
            if abs(amt - 0.0) < 0.00001:
                continue

            food = self.food_details[idx]

            cals = '?'*8
            for n in food["nutrients"]:
                if n["nutrient_id"] == "208":
                    cals = float(n["nutrient_val"]) * amt
                    total_cals += cals
                    cals = "%8.2f" % cals
                    break

            info("  FOOD %8.2f units (%s cals) of %s (grp %s)",
                amt,
                cals,
                food["short_descrip"],
                food["food_group_descrip"]
            )

        info("Total Calories: %8.2f", total_cals)

def food_scores(engine):
    """Just output all foods that we allow with the score that you would get."""
    one_hot = []
    for i in range(engine.food_count):
        one = [0.0] * engine.food_count
        one[i] = 1.0
        one_hot.append(list(one))

    engine.set_pop(one_hot)
    engine.generation()

    for food, nutr in zip(engine.last_population, engine.last_popnutrition):
        engine.dump_instance(food, nutr)
        info('='*78)

def main():
    input_file = "foods.jsontxt.gz"  # TODO: command line or stdin or...

    info("Micronutrients used for calcs: %d", len(ALL_MICROS))

    info("Reading input file %s", input_file)
    food_details = list(read_foods(input_file))

    engine = opt_engine(food_details)
    info("Init population shape: %s", engine.population.shape)
    info("Food matrix shape: %s", engine.foods.shape)

    if '-foodscores' in sys.argv[1:]:
        food_scores(engine)
        return
    elif '-topn' in sys.argv[1:]:
        print("TODO: top-n per nutrient output")
        return

    # Do our generations
    # TODO: good stopping criteria
    with open("results", "w") as results:
        def file_output(s, *args):
            if args:
                s = s % args
            results.write(s)
            results.write('\n')
            results.flush()

        info("Pre-Running 2 generations")
        engine.generation()
        engine.generation()

        seen = set()

        for gen in range(1, 5000+1):
            engine.generation()

            best_food, best_nutr = engine.last_population[0], engine.last_popnutrition[0]
            info("BEGIN Generation %d (pop size %d)", gen, len(engine.last_population))
            engine.dump_instance(best_food, best_nutr)
            info("END Generation %d", gen)
            info('-'*20)

            for food, nutr in zip(engine.last_population[:3], engine.last_popnutrition[:3]):
                key = tuple(food)
                if key in seen:
                    continue

                seen.add(key)

                file_output("First Seen Generation %d", gen)
                engine.dump_instance(food, nutr, info=file_output)
                file_output('='*78)

if __name__ == "__main__":
    main()


from torch import nn
import torch
import numpy as np

class SemiCRFShiftRelay(nn.Module):
    """
    该模块是一个decoder，但当前不支持含有tag的decode。

    """
    def __init__(self, L):
        """

        :param L: 不包含relay的长度
        """
        if L<2:
            raise RuntimeError()
        super().__init__()
        self.L = L

    def forward(self, logits, relay_logits, relay_target, relay_mask, end_seg_mask, seq_len):
        """
        relay node是接下来L个字都不是它的结束。relay的状态是往后滑动1个位置

        :param logits: batch_size x max_len x L, 当前位置往左边L个segment的分数，最后一维的0是长度为1的segment(即本身)
        :param relay_logits: batch_size x max_len, 当前位置是接下来L-1个位置都不是终点的分数
        :param relay_target: batch_size x max_len 每个位置他的segment在哪里开始的。如果超过L，则一直保持为L-1。比如长度为
            5的词，L=3, [0, 1, 2, 2, 2]
        :param relay_mask: batch_size x max_len, 在需要relay的地方为1, 长度为5的词, L=3时，为[1, 1, 1, 0, 0]
        :param end_seg_mask: batch_size x max_len, segment结束的地方为1。
        :param seq_len: batch_size, 句子的长度
        :return: loss: batch_size,
        """
        batch_size, max_len, L = logits.size()

        # 当前时刻为relay node的分数是多少
        relay_scores = logits.new_zeros(batch_size, max_len)
        # 当前时刻结束的分数是多少
        scores = logits.new_zeros(batch_size, max_len+1)
        # golden的分数
        gold_scores = relay_logits[:, 0].masked_fill(relay_mask[:, 0].eq(0), 0) + \
                                logits[:, 0, 0].masked_fill(end_seg_mask[:, 0].eq(0), 0)
        # 初始化
        scores[:, 1] = logits[:, 0, 0]
        batch_i = torch.arange(batch_size).to(logits.device).long()
        relay_scores[:, 0] = relay_logits[:, 0]
        last_relay_index = max_len - self.L
        for t in range(1, max_len):
            real_L = min(t+1, L)
            flip_logits_t = logits[:, t, :real_L].flip(dims=[1])  # flip之后低0个位置为real_L-1的segment
            # 计算relay_scores的更新
            if t<last_relay_index:
                #   (1) 从正常位置跳转
                tmp1 = relay_logits[:, t] + scores[:, t] # batch_size
                #   (2) 从relay跳转
                tmp2 = relay_logits[:, t] + relay_scores[:, t-1] # batch_size
                tmp1 = torch.stack([tmp1, tmp2], dim=0)
                relay_scores[:, t] = torch.logsumexp(tmp1, dim=0)
            # 计算scores的更新
            #  (1)从之前的位置跳转过来的
            tmp1 = scores[:, t-real_L+1:t+1] + flip_logits_t  # batch_size x L
            if t>self.L-1:
                #  (2)从relay跳转过来的
                tmp2 = relay_scores[:, t-self.L] # batch_size
                tmp2 = tmp2 + flip_logits_t[:, 0] # batch_size
                tmp1 = torch.cat([tmp1, tmp2.unsqueeze(-1)], dim=-1)
            scores[:, t+1] = torch.logsumexp(tmp1, dim=-1) # 更新当前时刻的分数

            # 计算golden
            seg_i = relay_target[:, t] # batch_size
            gold_segment_scores = logits[:, t][(batch_i, seg_i)].masked_fill(end_seg_mask[:, t].eq(0), 0) # batch_size, 后向从0到L长度的segment的分数
            relay_score = relay_logits[:, t].masked_fill(relay_mask[:, t].eq(0), 0)
            gold_scores = gold_scores + relay_score + gold_segment_scores
        all_scores = scores.gather(dim=1, index=seq_len.unsqueeze(1)).squeeze(1) # batch_size
        return all_scores - gold_scores

    def predict(self, logits, relay_logits, seq_len):
        """
        relay node是接下来L个字都不是它的结束。relay的状态是往后滑动L-1个位置

        :param logits: batch_size x max_len x L, 当前位置左边L个segment的分数，最后一维的0是长度为1的segment(即本身)
        :param relay_logits: batch_size x max_len, 当前位置是接下来L-1个位置都不是终点的分数
        :param seq_len: batch_size, 句子的长度
        :return: pred: batch_size x max_len以该点开始的segment的(长度-1); pred_mask为1的地方预测有segment开始
        """
        batch_size, max_len, L = logits.size()
        # 当前时刻为relay node的分数是多少
        max_relay_scores = logits.new_zeros(batch_size, max_len)
        relay_bt = seq_len.new_zeros(batch_size, max_len)  # 当前结果是否来自于relay的结果
        # 当前时刻结束的分数是多少
        max_scores = logits.new_zeros(batch_size, max_len+1)
        bt = seq_len.new_zeros(batch_size, max_len)
        # 初始化
        max_scores[:, 1] = logits[:, 0, 0]
        max_relay_scores[:, 0] = relay_logits[:, 0]
        last_relay_index = max_len - self.L
        for t in range(1, max_len):
            real_L = min(t+1, L)
            flip_logits_t = logits[:, t, :real_L].flip(dims=[1])  # flip之后低0个位置为real_L-1的segment
            # 计算relay_scores的更新
            if t<last_relay_index:
                #  (1) 从正常位置跳转
                tmp1 = relay_logits[:, t] + max_scores[:, t]
                #   (2) 从relay跳转
                tmp2 = relay_logits[:, t] + max_relay_scores[:, t-1] # batch_size
                # 每个sample的倒数L位不能是relay了
                tmp2 = tmp2.masked_fill(seq_len.le(t+L), float('-inf'))
                mask_i = tmp1.lt(tmp2) # 为1的位置为relay跳转
                relay_bt[:, t].masked_fill_(mask_i, 1)
                max_relay_scores[:, t] = torch.max(tmp1, tmp2)

            # 计算scores的更新
            #  (1)从之前的位置跳转过来的
            tmp1 = max_scores[:, t-real_L+1:t+1] + flip_logits_t  # batch_size x L
            tmp1 = tmp1.flip(dims=[1]) # 0的位置代表长度为1的segment
            if self.L-1<t:
                #  (2)从relay跳转过来的
                tmp2 = max_relay_scores[:, t-self.L] # batch_size
                tmp2 = tmp2 + flip_logits_t[:, 0]
                tmp1 = torch.cat([tmp1, tmp2.unsqueeze(-1)], dim=-1)
            # 看哪个更大
            max_score, pt = torch.max(tmp1, dim=1)
            max_scores[:, t+1] = max_score
            # mask_i = pt.ge(self.L)
            bt[:, t] = pt # 假设L=3, 那么对于0,1,2,3分别代表的是[t, t], [t-1, t], [t-2, t], [t-self.L(relay), t]
        # 需要把结果decode出来
        pred = np.zeros((batch_size, max_len), dtype=int)
        pred_mask = np.zeros((batch_size, max_len), dtype=int)
        seq_len = seq_len.tolist()
        bt = bt.tolist()
        relay_bt = relay_bt.tolist()
        for b in range(batch_size):
            seq_len_i = seq_len[b]
            bt_i = bt[b][:seq_len_i]
            relay_bt_i = relay_bt[b][:seq_len_i]
            j = seq_len_i - 1
            assert relay_bt_i[j]!=1
            while j>-1:
                if bt_i[j]==self.L:
                    seg_start_pos = j
                    j = j-self.L
                    while relay_bt_i[j]!=0 and j>-1:
                        j = j - 1
                    pred[b, j] = seg_start_pos - j
                    pred_mask[b, j] = 1
                else:
                    length = bt_i[j]
                    j = j - bt_i[j]
                    pred_mask[b, j] = 1
                    pred[b, j] = length
                j = j - 1

        return torch.LongTensor(pred).to(logits.device), torch.LongTensor(pred_mask).to(logits.device)

class FeatureFunMax(nn.Module):
    def __init__(self, hidden_size:int, L:int):
        """
        用于计算semi-CRF特征的函数。给定batch_size x max_len x hidden_size形状的输入，输出为batch_size x max_len x L的
        分数，以及batch_size x max_len的relay的分数。两者的区别参考论文 TODO 补充

        :param hidden_size: 输入特征的维度大小
        :param L: 不包含relay node的segment的长度大小。
        """
        super().__init__()

        self.end_fc = nn.Linear(hidden_size, 1, bias=False)
        self.whole_w = nn.Parameter(torch.randn(L, hidden_size))
        self.relay_fc = nn.Linear(hidden_size, 1)
        self.length_bias = nn.Parameter(torch.randn(L))
        self.L = L
    def forward(self, logits):
        """

        :param logits: batch_size x max_len x hidden_size
        :return: batch_size x max_len x L # 最后一维为左边segment的分数，0处为长度为1的segment
                 batch_size x max_len, # 当前位置是接下来L-1个位置都不是终点的分数

        """
        batch_size, max_len, hidden_size = logits.size()
        # start_scores = self.start_fc(logits) # batch_size x max_len x 1 # 每个位置作为start的分数
        tmp = logits.new_zeros(batch_size, max_len+self.L-1, hidden_size)
        tmp[:, -max_len:] = logits
        # batch_size x max_len x hidden_size x (self.L) -> batch_size x max_len x (self.L) x hidden_size
        start_logits = tmp.unfold(dimension=1, size=self.L, step=1).transpose(2, 3).flip(dims=[2])
        end_scores = self.end_fc(logits) # batch_size x max_len x 1
        # 计算relay的特征
        relay_tmp = logits.new_zeros(batch_size, max_len, hidden_size)
        relay_tmp[:, :-self.L] = logits[:, self.L:]
        # batch_size x max_len x hidden_size
        relay_logits_max = torch.max(relay_tmp, logits) # end - start
        logits_max = torch.max(logits.unsqueeze(2), start_logits) # batch_size x max_len x L x hidden_size
        whole_scores = (logits_max*self.whole_w).sum(dim=-1) # batch_size x max_len x self.L
        # whole_scores = self.whole_fc().squeeze(-1) # bz x max_len x self.L
        # batch_size x max_len
        relay_scores = self.relay_fc(relay_logits_max).squeeze(-1)
        return whole_scores+end_scores+self.length_bias.view(1, 1, -1), relay_scores


# In this tutorial, you will learn how to write functions in python
# Lets write a function that will return the square, square root and the cube of the number that is passed to it

import numpy as np

# This is how a function is defined:
def my_func(val): # the colon works the same way it did in for loops, the indenting begins after that (Do not delete the indenting)
    sq = val*val
    sqrt = np.sqrt(val)
    cube = val*val*val
    return sq,sqrt,cube

# you press shift+tab to come out of the indent (or loop)

v =100
print("The output from the first technique is:")
print(my_func(v))

# or you can also do:
val1, val2, val3 = my_func(v) # in this case, val1 contains sq, val2 contains sqrt and val3 contains cube
print("The output from the second technique is:")
print("value of val1 is {0}".format(val1))
print("value of val2 is {0}".format(val2))
print("value of val3 is {0}".format(val3))

# assignments

# write a function that will determine if a given number is even or odd
# write a function to find the sin of a number using taylors expansion
# write a function to find e^x of a number
# write a function that do all the above in a single go