"""
Utility functions and classes for 2D and 3D ocean models
"""
from __future__ import absolute_import

import os
import sys
from abc import ABCMeta, abstractmethod  # NOQA
from collections import OrderedDict, namedtuple  # NOQA

import coffee.base as ast  # NOQA
import ufl  # NOQA
from firedrake import *
from mpi4py import MPI  # NOQA
from pyop2.profiling import timed_function, timed_region, timed_stage  # NOQA

from .field_defs import field_metadata
from .log import *
from .physical_constants import physical_constants

ds_surf = ds_t
ds_bottom = ds_b

class FrozenClass(object):
    """A class where creating a new attribute will raise an exception if _isfrozen == True"""
    _isfrozen = False

    def __setattr__(self, key, value):
        if self._isfrozen and not hasattr(self, key):
            raise TypeError('Adding new attribute "{:}" to {:} class is forbidden'.format(key, self.__class__.__name__))
        super(FrozenClass, self).__setattr__(key, value)

class SumFunction(object):
    """
    Helper class to keep track of sum of Coefficients.
    """
    def __init__(self):
        """
        Initialize empty sum.

        get operation returns Constant(0)
        """
        self.coeff_list = []

    def add(self, coeff):
        """
        Adds a coefficient to self
        """
        if coeff is None:
            return
        self.coeff_list.append(coeff)

    def get_sum(self):
        """
        Returns a sum of all added Coefficients
        """
        if len(self.coeff_list) == 0:
            return None
        return sum(self.coeff_list)

class AttrDict(dict):
    """
    Dictionary that provides both self['key'] and self.key access to members.

    http://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute-in-python
    """
    def __init__(self, *args, **kwargs):
        if sys.version_info < (2, 7, 4):
            raise Exception('AttrDict requires python >= 2.7.4 to avoid memory leaks')
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self

class FieldDict(AttrDict):
    """
    AttrDict that checks that all added fields have proper meta data.

    Values can be either Function or Constant objects.
    """
    def _check_inputs(self, key, value):
        if key != '__dict__':
            from firedrake.functionspaceimpl import (MixedFunctionSpace,
                                                     WithGeometry)
            if not isinstance(value, (Function, Constant)):
                raise TypeError('Value must be a Function or Constant object')
            fs = value.function_space()
            is_mixed = (isinstance(fs, MixedFunctionSpace)
                        or (isinstance(fs, WithGeometry)
                            and isinstance(fs.topological, MixedFunctionSpace)))
            if not is_mixed and key not in field_metadata:
                msg = 'Trying to add a field "{:}" that has no metadata. ' \
                      'Add field_metadata entry to field_defs.py'.format(key)
                raise Exception(msg)

    def _set_functionname(self, key, value):
        """Set function.name to key to ensure consistent naming"""
        if isinstance(value, Function):
            value.rename(name=key)

    def __setitem__(self, key, value):
        self._check_inputs(key, value)
        self._set_functionname(key, value)
        super(FieldDict, self).__setitem__(key, value)

    def __setattr__(self, key, value):
        self._check_inputs(key, value)
        self._set_functionname(key, value)
        super(FieldDict, self).__setattr__(key, value)

def get_functionspace(mesh, h_family, h_degree, v_family=None, v_degree=None,
                      vector=False, hdiv=False, variant=None, v_variant=None,
                      **kwargs):
    cell_dim = mesh.cell_dimension()
    assert cell_dim in [2, (2, 1)], 'Unsupported cell dimension'
    hdiv_families = [
        'RT', 'RTF', 'RTCF', 'RAVIART-THOMAS',
        'BDM', 'BDMF', 'BDMCF', 'BREZZI-DOUGLAS-MARINI',
    ]
    if variant is None:
        if h_family.upper() in hdiv_families:
            if h_family in ['RTCF', 'BDMCF']:
                variant = 'equispaced'
            else:
                variant = 'integral'
        else:
            variant = 'equispaced'
    if v_variant is None:
        v_variant = 'equispaced'
    if cell_dim == (2, 1):
        if v_family is None:
            v_family = h_family
        if v_degree is None:
            v_degree = h_degree
        h_cell, v_cell = mesh.ufl_cell().sub_cells()
        h_elt = FiniteElement(h_family, h_cell, h_degree, variant=variant)
        v_elt = FiniteElement(v_family, v_cell, v_degree, variant=v_variant)
        elt = TensorProductElement(h_elt, v_elt)
        if hdiv:
            elt = HDiv(elt)
    else:
        elt = FiniteElement(h_family, mesh.ufl_cell(), h_degree, variant=variant)

    constructor = VectorFunctionSpace if vector else FunctionSpace
    return constructor(mesh, elt, **kwargs)

def get_extruded_base_element(ufl_element):
    """
    Return UFL TensorProductElement of an extruded UFL element.

    In case of a non-extruded mesh, returns the element itself.
    """
    if isinstance(ufl_element, ufl.HDivElement):
        ufl_element = ufl_element._element
    if isinstance(ufl_element, ufl.MixedElement):
        ufl_element = ufl_element.sub_elements()[0]
    if isinstance(ufl_element, ufl.VectorElement):
        ufl_element = ufl_element.sub_elements()[0]  # take the first component
    if isinstance(ufl_element, ufl.EnrichedElement):
        ufl_element = ufl_element._elements[0]
    return ufl_element

ElementContinuity = namedtuple("ElementContinuity", ["horizontal", "vertical"])
"""
A named tuple describing the continuity of an element in the horizontal/vertical direction.

The field value is one of "cg", "hdiv", or "dg".
"""

def element_continuity(ufl_element):
    """Return an :class:`ElementContinuity` instance with the
    continuity of a given element.

    :arg ufl_element: The UFL element to determine the continuity
        of.
    :returns: A new :class:`ElementContinuity` instance.
    """
    elem = ufl_element
    elem_types = {
        'Discontinuous Lagrange': 'dg',
        'Lagrange': 'cg',
        'Raviart-Thomas': 'hdiv',
        'RTCF': 'hdiv',
        'Brezzi-Douglas-Marini': 'hdiv',
        'BDMCF': 'hdiv',
        'Q': 'cg',
        'DQ': 'dg',
    }

    base_element = get_extruded_base_element(ufl_element)
    if isinstance(elem, ufl.HDivElement):
        horiz_type = 'hdiv'
        vert_type = 'hdiv'
    elif isinstance(base_element, ufl.TensorProductElement):
        a, b = base_element.sub_elements()
        horiz_type = elem_types[a.family()]
        vert_type = elem_types[b.family()]
    else:
        horiz_type = elem_types[base_element.family()]
        vert_type = horiz_type
    return ElementContinuity(horiz_type, vert_type)

def create_directory(path, comm=COMM_WORLD):
    """
    Create a directory on disk

    Raises IOError if a file with the same name already exists.
    """
    if comm.rank == 0:
        if os.path.exists(path):
            if not os.path.isdir(path):
                raise IOError('file with same name exists', path)
        else:
            os.makedirs(path)
    comm.barrier()
    return path

def get_facet_mask(function_space, facet='bottom'):
    """
    Returns the top/bottom nodes of extruded 3D elements.

    :arg function_space: Firedrake :class:`FunctionSpace` object
    :kwarg str facet: 'top' or 'bottom'

    .. note::
        The definition of top/bottom depends on the direction of the extrusion.
        Here we assume that the mesh has been extruded upwards (along positive
        z axis).
    """
    from tsfc.finatinterface import create_element as create_finat_element

    # get base element
    elem = get_extruded_base_element(function_space.ufl_element())
    assert isinstance(elem, TensorProductElement), \
        f'function space must be defined on an extruded 3D mesh: {elem}'
    # figure out number of nodes in sub elements
    h_elt, v_elt = elem.sub_elements()
    nb_nodes_h = create_finat_element(h_elt).space_dimension()
    nb_nodes_v = create_finat_element(v_elt).space_dimension()
    # compute top/bottom facet indices
    # extruded dimension is the inner loop in index
    # on interval elements, the end points are the first two dofs
    offset = 0 if facet == 'bottom' else 1
    indices = np.arange(nb_nodes_h)*nb_nodes_v + offset
    return indices

def extrude_mesh_sigma(mesh2d, n_layers, bathymetry_2d, z_stretch_fact=1.0,
                       min_depth=None):
    """
    Extrudes a 2d surface mesh with bathymetry data defined in a 2d field.

    Generates a uniform terrain following mesh.

    :arg mesh2d: 2D mesh
    :arg n_layers: number of vertical layers
    :arg bathymetry: 2D :class:`Function` of the bathymetry
        (the depth of the domain; positive downwards)
    """
    mesh = ExtrudedMesh(mesh2d, layers=n_layers, layer_height=1.0/n_layers)

    coordinates = mesh.coordinates
    fs_3d = coordinates.function_space()
    fs_2d = bathymetry_2d.function_space()
    new_coordinates = Function(fs_3d)

    z_stretch_func = Function(fs_2d)
    if isinstance(z_stretch_fact, Function):
        assert z_stretch_fact.function_space() == fs_2d
        z_stretch_func = z_stretch_fact
    else:
        z_stretch_func.assign(z_stretch_fact)

    # number of nodes in vertical direction
    n_vert_nodes = fs_3d.finat_element.space_dimension() / fs_2d.finat_element.space_dimension()

    min_depth_arr = np.ones((n_layers+1, ))*1e22
    if min_depth is not None:
        for i, v in enumerate(min_depth):
            min_depth_arr[i] = v

    nodes = get_facet_mask(fs_3d, 'bottom')
    idx = op2.Global(len(nodes), nodes, dtype=np.int32, name='node_idx')
    min_depth_op2 = op2.Global(len(min_depth_arr), min_depth_arr, name='min_depth')
    kernel = op2.Kernel("""
        void my_kernel(double *new_coords, double *old_coords, double *bath2d, double *z_stretch, int *idx, double *min_depth) {
            for ( int d = 0; d < %(nodes)d; d++ ) {
                double s_fact = z_stretch[d];
                for ( int e = 0; e < %(v_nodes)d; e++ ) {
                    new_coords[3*(idx[d]+e) + 0] = old_coords[3*(idx[d]+e) + 0];
                    new_coords[3*(idx[d]+e) + 1] = old_coords[3*(idx[d]+e) + 1];
                    double sigma = 1.0 - old_coords[3*(idx[d]+e) + 2]; // top 0, bot 1
                    double new_z = -bath2d[d] * pow(sigma, s_fact) ;
                    int layer = fmin(fmax(round(sigma*(%(n_layers)d + 1) - 1.0), 0.0), %(n_layers)d);
                    double max_z = -min_depth[layer];
                    new_z = fmax(new_z, max_z);
                    new_coords[3*(idx[d]+e) + 2] = new_z;
                }
            }
        }""" % {'nodes': fs_2d.finat_element.space_dimension(),
                'v_nodes': n_vert_nodes,
                'n_layers': n_layers},
        'my_kernel')

    op2.par_loop(kernel, mesh.cell_set,
                 new_coordinates.dat(op2.WRITE, fs_3d.cell_node_map()),
                 coordinates.dat(op2.READ, fs_3d.cell_node_map()),
                 bathymetry_2d.dat(op2.READ, fs_2d.cell_node_map()),
                 z_stretch_func.dat(op2.READ, fs_2d.cell_node_map()),
                 idx(op2.READ),
                 min_depth_op2(op2.READ),
                 iterate=op2.ALL)

    mesh.coordinates.assign(new_coordinates)

    return mesh

def comp_volume_2d(eta, bath):
    """Computes volume of the 2D domain as an integral of the elevation field"""
    val = assemble((eta+bath)*dx)
    return val

def comp_volume_3d(mesh):
    """Computes volume of the 3D domain as an integral"""
    one = Constant(1.0, domain=mesh.coordinates.ufl_domain())
    val = assemble(one*dx)
    return val

def comp_tracer_mass_2d(scalar_func, total_depth):
    """
    Computes total tracer mass in the 2D domain
    :arg scalar_func: depth-averaged scalar :class:`Function` to integrate
    :arg total_depth: scalar UFL expression (e.g. from get_total_depth())
    """
    val = assemble(scalar_func*total_depth*dx)
    return val

def comp_tracer_mass_3d(scalar_func):
    """
    Computes total tracer mass in the 3D domain

    :arg scalar_func: scalar :class:`Function` to integrate
    """
    val = assemble(scalar_func*dx)
    return val

def get_zcoord_from_mesh(zcoord, solver_parameters=None):
    """
    Evaluates z coordinates from the 3D mesh

    :arg zcoord: scalar :class:`Function` where coordinates will be stored
    """
    # TODO coordinates should probably be interpolated instead
    if solver_parameters is None:
        solver_parameters = {}
    solver_parameters.setdefault('ksp_atol', 1e-12)
    solver_parameters.setdefault('ksp_rtol', 1e-16)
    fs = zcoord.function_space()
    tri = TrialFunction(fs)
    test = TestFunction(fs)
    a = tri*test*dx
    l = fs.mesh().coordinates[2]*test*dx
    solve(a == l, zcoord, solver_parameters=solver_parameters)
    return zcoord

def compute_baroclinic_head(solver):
    r"""
    Computes the baroclinic head :math:`r` from the density field

    .. math::
        r = \frac{1}{\rho_0} \int_{z}^\eta  \rho' d\zeta.
    """
    with timed_stage('density_solve'):
        solver.density_solver.solve()
    with timed_stage('rho_integral'):
        solver.rho_integrator.solve()
        solver.fields.baroc_head_3d *= -physical_constants['rho0_inv']
    with timed_stage('int_pg_solve'):
        solver.int_pg_calculator.solve()

def extend_function_to_3d(func, mesh_extruded):
    """
    Returns a 3D view of a 2D :class:`Function` on the extruded domain.

    The 3D function resides in V x R function space, where V is the function
    space of the source function. The 3D function shares the data of the 2D
    function.
    """
    fs = func.function_space()
    assert fs.mesh().geometric_dimension() == 2, 'Function must be in 2D space'
    ufl_elem = fs.ufl_element()
    family = ufl_elem.family()
    degree = ufl_elem.degree()
    name = func.name()
    if isinstance(ufl_elem, ufl.VectorElement):
        # vector function space
        fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0,
                                        dim=2, vector=True)
    else:
        fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0)
    func_extended = Function(fs_extended, name=name, val=func.dat._data)
    func_extended.source = func
    return func_extended

class ExtrudedFunction(Function):
    """
    A 2D :class:`Function` that provides a 3D view on the extruded domain.

    The 3D function can be accessed as `ExtrudedFunction.view_3d`.
    The 3D function resides in V x R function space, where V is the function
    space of the source function. The 3D function shares the data of the 2D
    function.
    """
    def __init__(self, *args, mesh_3d=None, **kwargs):
        """
        Create a 2D :class:`Function` with a 3D view on extruded mesh.

        :arg mesh_3d: Extruded 3D mesh where the function will be extended to.
        """
        # create the 2d function
        super().__init__(*args, **kwargs)

        if mesh_3d is not None:
            self.view_3d = extend_function_to_3d(self, mesh_3d)

class SubdomainProjector(object):
    """
    Projector that projects the restriction of an expression to the specified subdomain.
    """
    def __init__(self, v, v_out, subdomain_id, solver_parameters=None, constant_jacobian=True):

        if not isinstance(v, (ufl.core.expr.Expr, Function)):
            raise ValueError("Can only project UFL expression or Functions not '%s'" % type(v))

        self.v = v
        self.v_out = v_out

        V = v_out.function_space()

        p = TestFunction(V)
        q = TrialFunction(V)

        a = inner(p, q)*dx
        L = inner(p, v)*dx(subdomain_id)

        problem = LinearVariationalProblem(a, L, v_out,
                                           constant_jacobian=constant_jacobian)

        if solver_parameters is None:
            solver_parameters = {}

        solver_parameters.setdefault("ksp_type", "cg")

        self.solver = LinearVariationalSolver(problem,
                                              solver_parameters=solver_parameters)

    def project(self):
        """
        Apply the projection.
        """
        self.solver.solve()

def compute_elem_height(zcoord, output):
    """
    Computes the element height on an extruded mesh.

    :arg zcoord: field that contains the z coordinates of the mesh
    :type zcoord: :class:`Function`
    :arg output: field where element height is stored
    :type output: :class:`Function`
    """
    fs_in = zcoord.function_space()
    fs_out = output.function_space()

    iterate = op2.ALL

    # NOTE height maybe <0 if mesh was extruded like that
    kernel = op2.Kernel("""
        void my_kernel(double *func, double *zcoord) {
            for ( int d = 0; d < %(nodes)d/2; d++ ) {
                for ( int c = 0; c < %(func_dim)d; c++ ) {
                    double dz = fabs(zcoord[%(func_dim)d*(2*d+1) + c] - zcoord[%(func_dim)d*2*d + c]);
                    func[%(output_dim)d*2*d + c] = dz;
                    func[%(output_dim)d*(2*d+1) + c] = dz;
                }
            }
        }""" % {'nodes': zcoord.cell_node_map().arity,
                'func_dim': zcoord.function_space().value_size,
                'output_dim': output.function_space().value_size},
        'my_kernel')
    op2.par_loop(
        kernel, fs_out.mesh().cell_set,
        output.dat(op2.WRITE, fs_out.cell_node_map()),
        zcoord.dat(op2.READ, fs_in.cell_node_map()),
        iterate=iterate)

    return output

def get_horizontal_elem_size_2d(sol2d):
    """
    Computes horizontal element size from the 2D mesh

    :arg sol2d: 2D :class:`Function` where result is stored
    """
    p1_2d = sol2d.function_space()
    mesh = p1_2d.mesh()
    test = TestFunction(p1_2d)
    tri = TrialFunction(p1_2d)
    a = inner(test, tri) * dx
    l = inner(test, sqrt(CellVolume(mesh))) * dx
    solve(a == l, sol2d)

def get_facet_areas(mesh):
    """
    Compute area of each facet of `mesh`. The facet areas are stored as a HDiv trace field.

    NOTES:
      * In the 2D case, this gives edge lengths.
      * The plus sign is arbitrary and could equally well be chosen as minus.
    """
    HDivTrace = FunctionSpace(mesh, "HDiv Trace", 0)
    v, u = TestFunction(HDivTrace), TrialFunction(HDivTrace)
    facet_areas = Function(HDivTrace, name="Facet areas")
    mass_term = v('+')*u('+')*dS + v*u*ds
    rhs = v('+')*FacetArea(mesh)*dS + v*FacetArea(mesh)*ds
    solve(mass_term == rhs, facet_areas)
    return facet_areas

def get_minimum_angles_2d(mesh2d):
    """
    Compute the minimum angle in each element of a triangular mesh, `mesh2d`, using the
    cosine rule. The minimum angles are outputted as a P0 field.
    """
    try:
        assert mesh2d.topological_dimension() == 2
        assert mesh2d.ufl_cell() == ufl.triangle
    except AssertionError:
        raise NotImplementedError("Minimum angle only currently implemented for triangles.")
    edge_lengths = get_facet_areas(mesh2d)
    min_angles = Function(FunctionSpace(mesh2d, "DG", 0))
    par_loop("""for (int i=0; i<angle.dofs; i++) {

                  double min_edge = edges[0];
                  int min_index = 0;

                  for (int j=1; j<3; j++){
                    if (edges[j] < min_edge) {
                      min_edge = edges[j];
                      min_index = j;
                    }
                  }

                  double numerator = 0.0;
                  double denominator = 2.0;

                  for (int j=0; j<3; j++){
                    if (j == min_index) {
                      numerator -= edges[j]*edges[j];
                    } else {
                      numerator += edges[j]*edges[j];
                      denominator *= edges[j];
                    }
                  }
                  angle[0] = acos(numerator/denominator);
                }""", dx, {'edges': (edge_lengths, READ), 'angle': (min_angles, RW)})
    return min_angles

def get_cell_widths_2d(mesh2d):
    """
    Compute widths of mesh elements in each coordinate direction as the maximum distance
    between components of vertex coordinates.
    """
    try:
        assert mesh2d.topological_dimension() == 2
        assert mesh2d.ufl_cell() == ufl.triangle
    except AssertionError:
        raise NotImplementedError("Cell widths only currently implemented for triangles.")
    cell_widths = Function(VectorFunctionSpace(mesh2d, "DG", 0)).assign(np.finfo(0.0).min)
    par_loop("""for (int i=0; i<coords.dofs; i++) {
                  widths[0] = fmax(widths[0], fabs(coords[2*i] - coords[(2*i+2)%6]));
                  widths[1] = fmax(widths[1], fabs(coords[2*i+1] - coords[(2*i+3)%6]));
                }""", dx, {'coords': (mesh2d.coordinates, READ), 'widths': (cell_widths, RW)})
    return cell_widths

def anisotropic_cell_size(mesh):
    """
    Measure of cell size for anisotropic meshes, as described in
    Micheletti et al. (2003).

    This is used in the SUPG formulation for the 2D tracer model.

    Micheletti, Perotto and Picasso (2003). Stabilized finite
    elements on anisotropic meshes: a priori error estimates for
    the advection-diffusion and the Stokes problems. SIAM Journal
    on Numerical Analysis 41.3: 1131-1162.
    """
    try:
        from firedrake.slate.slac.compiler import PETSC_ARCH
    except ImportError:
        PETSC_ARCH = os.path.join(os.environ.get('PETSC_DIR'), os.environ.get('PETSC_ARCH'))
    include_dir = ["%s/include/eigen3" % PETSC_ARCH]

    # Compute cell Jacobian
    P0_ten = TensorFunctionSpace(mesh, "DG", 0)
    J = Function(P0_ten, name="Cell Jacobian")
    J.interpolate(Jacobian(mesh))

    # Compute minimum eigenvalue
    P0 = FunctionSpace(mesh, "DG", 0)
    min_evalue = Function(P0, name="Minimum eigenvalue")
    kernel_str = """
#include <Eigen/Dense>

using namespace Eigen;

void eigmin(double minEval[1], const double * J_) {
  Matrix<double, 2, 2, RowMajor> A;

  // Map input onto an Eigen object
  Map<Matrix<double, 2, 2, RowMajor> > J((double *)J_);

  // Compute singular value decomposition
  JacobiSVD<Matrix<double, 2, 2, RowMajor> > svd(J, ComputeFullV);

  // Get SPD part of polar decomposition
  A = svd.matrixV() * svd.singularValues().asDiagonal() * svd.matrixV().transpose();

  // Solve eigenvalue problem
  SelfAdjointEigenSolver<Matrix<double, 2, 2, RowMajor>> eigensolver(A);
  Vector2d D = eigensolver.eigenvalues();

  // Select minimum eigenvalue in modulus
  if (fabs(D(0)) < fabs(D(1))) minEval[0] = D(0);
  else minEval[0] = D(1);
}
"""
    kernel = op2.Kernel(kernel_str, 'eigmin', cpp=True, include_dirs=include_dir)
    op2.par_loop(kernel, P0_ten.node_set, min_evalue.dat(op2.RW), J.dat(op2.READ))
    return min_evalue

def beta_plane_coriolis_params(latitude):
    r"""
    Computes beta plane parameters :math:`f_0,\beta` based on latitude

    :arg float latitude: latitude in degrees
    :return: f_0, beta
    :rtype: float
    """
    omega = 7.2921150e-5  # rad/s Earth rotation rate
    r = 6371.e3  # Earth radius
    # Coriolis parameter f = 2 Omega sin(alpha)
    # Beta plane approximation f_beta = f_0 + Beta y
    # f_0 = 2 Omega sin(alpha_0)
    # Beta = df/dy|_{alpha=alpha_0}
    #      = (df/dalpha*dalpha/dy)_{alpha=alpha_0}
    #      = 2 Omega cos(alpha_0) /R
    alpha_0 = 2*np.pi*latitude/360.0
    f_0 = 2*omega*np.sin(alpha_0)
    beta = 2*omega*np.cos(alpha_0)/r
    return f_0, beta

def beta_plane_coriolis_function(latitude, out_function, y_offset=0.0):
    """
    Interpolates beta plane Coriolis function to a field

    :arg float latitude: latitude in degrees
    :arg out_function: :class:`Function` where to interpolate
    :kwarg float y_offset: offset (y - y_0) used in Beta-plane approximation.
        A constant in mesh coordinates.
    """
    # NOTE assumes that mesh y coordinate spans [-L_y, L_y]
    f0, beta = beta_plane_coriolis_params(latitude)
    coords = SpatialCoordinate(out_function.function_space().mesh())
    out_function.interpolate(f0 + beta * (coords[1] - y_offset))

def tensor_jump(v, n):
    r"""
    Jump term for vector functions based on the tensor product

    .. math::
        \text{jump}(\mathbf{u}, \mathbf{n}) = (\mathbf{u}^+ \mathbf{n}^+) +
        (\mathbf{u}^- \mathbf{n}^-)

    This is the discrete equivalent of grad(u) as opposed to the
    vectorial UFL jump operator :meth:`ufl.jump` which represents div(u).
    """
    return outer(v('+'), n('+')) + outer(v('-'), n('-'))

def compute_boundary_length(mesh2d):
    """
    Computes the length of the boundary segments in given 2d mesh
    """
    p1 = get_functionspace(mesh2d, 'CG', 1)
    boundary_markers = sorted(mesh2d.exterior_facets.unique_markers)
    boundary_len = OrderedDict()
    for i in boundary_markers:
        ds_restricted = ds(int(i))
        one_func = Function(p1).assign(1.0)
        boundary_len[i] = assemble(one_func * ds_restricted)
    return boundary_len

def select_and_move_detectors(mesh, detector_locations, detector_names=None,
                              maximum_distance=0.):
    """Select those detectors that are within the domain and/or move them to
    the nearest cell centre within the domain

    :arg mesh: Defines the domain in which detectors are to be located
    :arg detector_locations: List of x, y locations
    :arg detector_names: List of detector names (optional). If provided, a list
       of selected locations and a list of selected detector names are returned,
       otherwise only a list of selected locations is returned
    :arg maximum_distance: Detectors whose initial locations is outside the domain,
      but for which the nearest cell centre is within the specified distance, are
      moved to this location. By default a maximum distance of 0.0 is used, i.e
      no detectors are moved.
    """
    # auxilary function to test whether we can interpolate it in the given locations
    V = FunctionSpace(mesh, "CG", 1)
    v = Function(V)

    P0 = FunctionSpace(mesh, "DG", 0)
    VP0 = VectorFunctionSpace(mesh, "DG", 0)
    dist = Function(P0)
    loc_const = Constant(detector_locations[0])
    xy = SpatialCoordinate(mesh)
    p0xy = Function(VP0).interpolate(xy)

    # comparison operator that sorts on first entry first, etc.
    def min_lexsort(x, y, datatype):
        for xi, yi in zip(x, y):
            if xi < yi:
                return x
            elif yi < xi:
                return y
        # all entries the same:
        return x
    min_lexsort_op = MPI.Op.Create(min_lexsort, commute=False)

    def move_to_nearest_cell_center(location):
        loc_const.assign(location)
        dist.interpolate(dot(xy-loc_const, xy-loc_const))
        ind = dist.dat.data_ro.argmin()
        # smallest distance to a cell centre location on this process:
        local_loc = list(p0xy.dat.data_ro[ind])
        local_dist = np.sqrt(dist.dat.data_ro[ind])
        # select the smallest distance on all processes. If some distances are equal, pick a unique loc. based on lexsort
        global_dist_loc = mesh.comm.allreduce([local_dist]+local_loc, op=min_lexsort_op)
        return global_dist_loc[0], global_dist_loc[1:]

    accepted_locations = []
    accepted_names = []
    if detector_names is None:
        names = [None] * len(detector_locations)
    else:
        names = detector_names
    for location, name in zip(detector_locations, names):
        try:
            v(location)
        except PointNotInDomainError:
            moved_dist, location = move_to_nearest_cell_center(location)
            if moved_dist > maximum_distance:
                continue
        accepted_locations.append(location)
        accepted_names.append(name)

    min_lexsort_op.Free()

    if detector_names is None:
        return accepted_locations
    else:
        return accepted_locations, accepted_names

class DepthExpression:
    r"""
    Construct expression for depth depending on options

    If `not use_nonlinear_equations`, then the depth is simply the bathymetry:

  ..math::
        H = h

    Otherwise we include the free surface elevation:

  ..math::
        H = h + \eta

    and if `use_wetting_and_drying`, includes a bathymetry displacement term
    to ensure a positive depth (see Karna et al. 2011):

  ..math::
        H = h + f(h+\eta) + \eta

    where

  ..math::
        f(h+\eta) = (\sqrt{(h+\eta)^2 +\alpha^2} - (h+\eta))/2

    This introduces a wetting-drying parameter :math:`\alpha`, with dimensions
    of length. The value for :math:`\alpha` is specified by
    :attr:`.ModelOptions.wetting_and_drying_alpha`, in units of meters. The
    default value is 0.5, but the appropriate value is problem specific and
    should be set by the user.
    """

    def __init__(self, bathymetry_2d, use_nonlinear_equations=True,
                 use_wetting_and_drying=False, wetting_and_drying_alpha=0.5):
        self.bathymetry_2d = bathymetry_2d
        self.use_nonlinear_equations = use_nonlinear_equations
        self.use_wetting_and_drying = use_wetting_and_drying
        self.wetting_and_drying_alpha = wetting_and_drying_alpha

    def wd_bathymetry_displacement(self, eta):
        """
        Returns wetting and drying bathymetry displacement as described in:
        Karna et al.,  2011.
        :arg eta: current elevation as UFL expression
        """
        if self.use_wetting_and_drying:
            H = self.bathymetry_2d + eta
            return 0.5 * (sqrt(H ** 2 + self.wetting_and_drying_alpha ** 2) - H)
        else:
            return 0

    def get_total_depth(self, eta):
        """
        Returns total water column depth based on options
        :arg eta: current elevation as UFL expression
        """
        if self.use_nonlinear_equations:
            total_h = self.bathymetry_2d + eta + self.wd_bathymetry_displacement(eta)
        else:
            total_h = self.bathymetry_2d
        return total_h

class VorticityCalculator2D(object):
    r"""
    Linear solver for recovering fluid vorticity.

    It is recommended that the vorticity is sought
    in :math:`\mathbb P1` space.

    :arg vorticity_2d: :class:`Function` to hold calculated vorticity.
    :arg solver_obj: :class:`FlowSolver2d` object.
    :kwargs: to be passed to the :class:`LinearVariationalSolver`.
    """
    def __init__(self, vorticity_2d, solver_obj, **kwargs):
        fs = vorticity_2d.function_space()
        if element_continuity(fs.ufl_element()).horizontal != 'cg':
            raise NotImplementedError
        n = FacetNormal(fs.mesh())

        # Weak formulation
        test = TestFunction(fs)
        uv, elev = split(solver_obj.fields.solution_2d)
        a = TrialFunction(fs)*test*dx
        L = -inner(perp(uv), grad(test))*dx \
            + dot(perp(uv), n)*test*ds

        # Setup vorticity solver
        prob = LinearVariationalProblem(a, L, vorticity_2d)
        sp = {
            'ksp_type': 'gmres',
            'ksp_gmres_restart': 20,
            'ksp_rtol': 1.0e-05,
            'pc_type': 'sor',
        }
        kwargs.setdefault('solver_parameters', sp)
        self.solver = LinearVariationalSolver(prob, **kwargs)

        if 'vorticity_2d' not in solver_obj.options.fields_to_export:
            print_output("NOTE: 'vorticity_2d' will be calculated but not exported.")

    def solve(self):
        self.solver.solve()

class DepthIntegratedPoissonSolver(object):
    r"""
    Construct solvers for Poisson equation and updating velocities

    Poisson equation is related to 2d NH SWE system

    Non-hydrostatic pressure :math:`q` is obtained from the generic form of equation

    .. math::
        \nabla \cdot \nabla q^{n+1/2} + A \cdot \nabla q^{n+1/2} + B q^{n+1/2} + C = 0

    The parameter terms A, B and C are defined as

    .. math::
        A = \frac{\nabla (\eta^* - d)}{H^*}
        B = \nabla A - \frac{4}{(H^*)^2}
        C = -2 \frac{\rho_0}{\Delta t} ( \nabla \cdot \bar{\textbf{u}}^* + 2 \frac{\bar{w} - w_b}{H^*} )

    where the :math:`H = \eta + d` denotes the water depth
    and the superscript star symbol represents the intermediate level of terms.
    """
    def __init__(self, q_2d, uv_2d, w_2d, elev_2d, depth, dt, bnd_functions=None, solver_parameters=None):
        if solver_parameters is None:
            solver_parameters = {'snes_type': 'ksponly',
                                 'ksp_type': 'preonly',
                                 'mat_type': 'aij',
                                 'pc_type': 'lu'}
        rho_0 = physical_constants['rho0']
        self.q_2d = q_2d
        self.uv_2d = uv_2d
        self.w_2d = w_2d
        self.elev_2d = elev_2d
        self.depth = depth
        self.dt = dt
        self.bnd_functions = bnd_functions

        fs_q = self.q_2d.function_space()
        test_q = TestFunction(fs_q)
        normal = FacetNormal(fs_q.mesh())
        boundary_markers = fs_q.mesh().exterior_facets.unique_markers

        bath_2d = self.depth.bathymetry_2d
        h_star = self.depth.get_total_depth(self.elev_2d)
        w_b = -dot(self.uv_2d, grad(bath_2d))  # TODO account for bed movement

        # weak form of `div(grad(q^{n+1/2}))`
        f = -dot(grad(self.q_2d), grad(test_q))*dx
        # weak form of `dot(grad(self.elev_2d - bath_2d)/h_star, grad(q^{n+1/2}))`
        grad_hori = grad(self.elev_2d - bath_2d)
        f += dot(grad_hori/h_star, grad(self.q_2d))*test_q*dx
        # weak form of `q^{n+1/2}*div(grad(self.elev_2d - bath_2d))/h_star`
        f += -dot(grad(self.q_2d*test_q/h_star), grad_hori)*dx
        # weak form of `-q^{n+1/2}*(dot(grad(self.elev_2d - bath_2d), grad(h_star)) + 4)/h_star**2`
        f += -(dot(grad_hori, grad(h_star)) + 4.)/h_star**2*self.q_2d*test_q*dx
        # weak form of `-2.*rho_0/self.dt*(div(self.uv_2d) + (self.w_2d - w_b)/(0.5*h_star))`
        const = 2.*rho_0/self.dt
        f += const*(dot(grad(test_q), self.uv_2d)*dx - 2.*(self.w_2d - w_b)/h_star*test_q*dx)

        # boundary conditions
        bcs = []
        for bnd_marker in boundary_markers:
            func = self.bnd_functions['shallow_water'].get(bnd_marker)
            ds_bnd = ds(int(bnd_marker))
            if func is not None:  # e.g. inlet flow, TODO be more precise
                bc = DirichletBC(fs_q, 0., int(bnd_marker))
                bcs.append(bc)
                f += self.q_2d*test_q/h_star*dot(grad_hori, normal)*ds_bnd
                f += -const*dot(self.uv_2d, normal)*test_q*ds_bnd

        prob_q = NonlinearVariationalProblem(f, self.q_2d)
        self.solver_q = NonlinearVariationalSolver(
            prob_q,
            solver_parameters=solver_parameters,
            options_prefix='poisson_solver'
        )
        # horizontal velocity updater
        fs_u = self.uv_2d.function_space()
        tri_u = TrialFunction(fs_u)
        test_u = TestFunction(fs_u)
        a_u = inner(tri_u, test_u)*dx
        l_u = dot(self.uv_2d - 0.5*self.dt/rho_0*(grad(self.q_2d) + grad_hori/h_star*self.q_2d), test_u)*dx
        prob_u = LinearVariationalProblem(a_u, l_u, self.uv_2d)
        self.solver_u = LinearVariationalSolver(prob_u)
        # vertical velocity updater
        fs_w = self.w_2d.function_space()
        tri_w = TrialFunction(fs_w)
        test_w = TestFunction(fs_w)
        a_w = inner(tri_w, test_w)*dx
        l_w = dot(self.w_2d + self.dt/rho_0*(self.q_2d/h_star), test_w)*dx
        prob_w = LinearVariationalProblem(a_w, l_w, self.w_2d)
        self.solver_w = LinearVariationalSolver(prob_w)

    def solve(self, solve_w=True):
        # solve non-hydrostatic pressure q
        self.solver_q.solve()
        # update horizontal velocity uv_2d
        self.solver_u.solve()
        # update vertical velocity w_2d
        if solve_w:
            self.solver_w.solve()


from django.shortcuts import render
from django.contrib.auth.decorators import login_required
from django.utils.decorators import method_decorator
from django.views.generic import ListView
from django.urls import reverse
from django.http import HttpResponseRedirect, Http404
from django.forms import inlineformset_factory, forms
from .models import Gun, Bullet, TestResult, Velocity
from .forms import GunForm, BulletForm, ResultForm, VelocityForm

class HomeView(ListView):
    model = Gun
    template_name = "home.html"

class GunsView(ListView):
    model = Gun
    template_name = "guns.html"

@login_required
def gunview(request, gun_id):
    """Show a single gun and all its info."""
    guns = Gun.objects.filter(owner=request.user, id=gun_id).prefetch_related('bullets__results__velocity')
    # Make sure the gun belongs to the current user.
    for gun in guns:
        if gun.owner != request.user:
            raise Http404

    context = {'guns': guns}
    return render(request, 'gun.html', context)

@login_required
def AddGun(request):
    """Add a new gun."""
    if request.method != 'POST':
        # No data submitted; create a blank form.
        form = GunForm()
    else:
        # POST data submitted; process data.
        form = GunForm(request.POST)
        if form.is_valid():
            new_gun = form.save(commit=False)
            new_gun.owner = request.user
            form.save()
            return HttpResponseRedirect(reverse('guns'))
    context = {'form': form}
    return render(request, 'add_gun.html', context)

@login_required
def addbullet(request, gun_id):
    """Add a new bullet."""
    gun = Gun.objects.get(id=gun_id)
    if request.method != 'POST':
        # No data submitted; create a blank form.
        form = BulletForm()
    else:
        # POST data submitted; process data.
        form = BulletForm(data=request.POST)
        if form.is_valid():
            new_bullet = form.save(commit=False)
            new_bullet.gun = gun
            new_bullet.save()
            return HttpResponseRedirect(reverse('gun', args=[gun_id]))
    context = {'form': form}
    return render(request, 'add_bullet.html', context)

@login_required
def add_data(request, bullet_id):
    """Add a new result."""
    ResultFormSet = inlineformset_factory(Bullet, TestResult, fields=('charge', 'moa'), can_delete=False, extra=5)
    bullet = Bullet.objects.get(id=bullet_id)
    if request.method != 'POST':
        # No data submitted; create a blank form.
        formset = ResultFormSet()
    else:
        # POST data submitted; process data.
        formset = ResultFormSet(request.POST, instance=bullet)
        if formset.is_valid():
            formset.save()
            return HttpResponseRedirect(reverse('gun', args=[bullet.gun.pk]))
    context = {'formset': formset}
    return render(request, 'add_data.html', context)

@login_required
def edit_data(request, result_id):
    """Edit an existing result."""
    result = TestResult.objects.get(id=result_id)
    gun = result.bullet.gun
    if request.method != 'POST':
        # Initial request; pre-fill form with the current entry.
        form = ResultForm(instance=result)
    else:
        # POST data submitted; process data.
        form = ResultForm(instance=result, data=request.POST)
        if form.is_valid():
            form.save()
            return HttpResponseRedirect(reverse('gun', args=[gun.id]))
    context = {'result': result, 'gun': gun, 'form': form}
    return render(request, 'edit_data.html', context)

@login_required
def delete_data(request, result_id):
    """delete an existing result."""
    result = TestResult.objects.get(id=result_id)
    result.delete()
    gun = result.bullet.gun
    return HttpResponseRedirect(reverse('gun', args=[gun.id]))

@login_required
def add_velocity(request, result_id):
    """Add new velocity."""
    VelocityFormSet = inlineformset_factory(TestResult, Velocity, fields=('shotnumber', 'velocity'), can_delete=False)
    result = TestResult.objects.get(id=result_id)
    if request.method != 'POST':
        # No data submitted; create a blank form.
        formset = VelocityFormSet(instance=result)
    else:
        # POST data submitted; process data.
        formset = VelocityFormSet(request.POST, instance=result)
        if formset.is_valid():
            formset.save()
            return HttpResponseRedirect(reverse('edit_data', args=[result.pk]))
    context = {'formset': formset}
    return render(request, 'add_velocity.html', context)

@login_required
def edit_bullet(request, bullet_id):
    """Edit an existing result."""
    bullet = Bullet.objects.get(id=bullet_id)
    gun = bullet.gun
    if request.method != 'POST':
        # Initial request; pre-fill form with the current entry.
        form = BulletForm(instance=bullet)
    else:
        # POST data submitted; process data.
        form = BulletForm(instance=bullet, data=request.POST)
        if form.is_valid():
            form.save()
            return HttpResponseRedirect(reverse('gun', args=[gun.id]))
    context = {'bullet': bullet, 'gun': gun, 'form': form}
    return render(request, 'edit_bullet.html', context)

@login_required
def delete_bullet(request, bullet_id):
    """delete an existing result."""
    bullet = Bullet.objects.get(id=bullet_id)
    bullet.delete()
    gun = bullet.gun
    return HttpResponseRedirect(reverse('gun', args=[gun.id]))

@login_required
def view_graph(request, bullet_id):
    """Show a single bullet graph."""
    bullet = Bullet.objects.filter(id=bullet_id).prefetch_related('results__velocity')
    velocitylist = []
    chargelist = []
    moalist = []
    for b in bullet:
        for result in b.results.all():
            chargelist.append(result.charge)
            moalist.append(result.moa)
            total_avg = 0
            for velocity in result.velocity.all():
                total_avg += velocity.velocity
            total_avg /= 3
            velocitylist.append(total_avg)

    context = {'bullet': bullet, 'velocitylist': velocitylist, 'chargelist':chargelist, 'moalist':moalist}
    return render(request, 'graph.html', context)


"""
# mkb_sammlung_europa.py
This DAG updates the following datasets:

- [100148](https://data.bs.ch/explore/dataset/100148)

"""
from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime, timedelta
from airflow.operators.docker_operator import DockerOperator

default_args = {
    'owner'                 : 'hester.pieters',
    'description'           : 'Run the mkb_sammlung_europa.py docker container',
    'depend_on_past'        : False,
    'start_date'            : datetime(2022, 1, 27),
    'email'                 : ["jonas.bieri@bs.ch", "jonas.eckenfels@bs.ch", "orhan.saeedi@bs.ch", "nicolas.maire@bs.ch"],
    'email_on_failure'      : True,
    'email_on_retry'        : False,
    'retries'               : 0,
    'retry_delay'           : timedelta(minutes=15)
}

with DAG('mkb_sammlung_europa', default_args=default_args, schedule_interval=None, catchup=False) as dag:
    dag.doc_md = __doc__
    upload = DockerOperator(
        task_id='upload',
        image='mkb_sammlung_europa:latest',
        api_version='auto',
        auto_remove=True,
        command='/bin/bash /code/data-processing/mkb_sammlung_europa/etl.sh ',
        container_name='mkb_sammlung_europa--upload',
        docker_url="unix://var/run/docker.sock",
        network_mode="bridge",
        tty=True,
        volumes=['/data/dev/workspace/data-processing:/code/data-processing', '/mnt/OGD-DataExch/PD-Kultur-MKB:/code/data-processing/mkb_sammlung_europa/data']
    )


from typing import List

class Solution:
    def restoreString(self, s: str, indices: List[int]) -> str:
        ans = [""] * len(s)

        for i in range(len(s)):
            ans[indices[i]] = s[i]
        return "".join(ans)

assert Solution().restoreString("codeleet", [4, 5, 6, 7, 0, 2, 1, 3])


import torchvision

from torchvision import transforms

def get_dataset(dataset_name, datadir):
    if dataset_name == "mnist":
        dataset = torchvision.datasets.MNIST(datadir, train=True,
                               download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.5,), (0.5,))
                               ]))

    if dataset_name == "cifar10":
        transform_function = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465),
                                 (0.2023, 0.1994, 0.2010)),
        ])

        dataset = torchvision.datasets.CIFAR10(
            root=datadir,
            train=True,
            download=True,
            transform=transform_function)

    if dataset_name == "cifar100":
        transform_function = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465),
                                 (0.2023, 0.1994, 0.2010)),
        ])

        dataset = torchvision.datasets.CIFAR100(
            root=datadir,
            train=True,
            download=True,
            transform=transform_function)

    return dataset


import filecmp
import json
import logging
import os
import random
import re
import tempfile
from collections import namedtuple
from itertools import combinations
from pathlib import Path
from typing import Callable, Optional, Set

import boto3
from tqdm import tqdm

from .util import Oracle, Results, instances, load_json

bucket = 'gnosis-dev-dfusion'
page_size = 1000

def get_instance_list(
    client,
    prefix: str
):
    """Return set of filename is S3 bucket."""
    l = set()
    page = client.list_objects_v2(
        Bucket=bucket, Prefix=prefix, MaxKeys=page_size
    )
    l |= {r['Key'] for r in page['Contents'] if r['Key'][-4:] == 'json'}

    while page['IsTruncated']:
        page = client.list_objects_v2(
            Bucket=bucket,
            Prefix=prefix,
            MaxKeys=page_size,
            ContinuationToken=page['NextContinuationToken']
        )
        l |= {r['Key'] for r in page['Contents'] if r['Key'][-4:] == 'json'}
    return l

def get_instance_date(instance_path: str):
    return re.search(r'\/([0-9\-]+)', instance_path).group(1)

def get_instance_batch(instance_path: str):
    return re.search(r'instance_([0-9]+)', instance_path).group(1)

def filter_instances_by_date(instance_paths, from_date):
    instance_paths = {
        instance_path
        for instance_path in instance_paths
        if get_instance_date(instance_path) >= from_date
    }
    return instance_paths

def filter_instances_by_batch(instance_paths, batches):
    instance_paths = {
        instance_path
        for instance_path in instance_paths
        if get_instance_batch(instance_path) in batches
    }
    return instance_paths

def get_instance_name(instance_path: str):
    return Path(instance_path).name

def get_result_name(result_path: str):
    return Path(result_path).parent.name

def download_instance(
    client,
    path: str,
    output_dir: Path,
    name_extractor: Callable[[str], str]
):
    instance_name = name_extractor(path)
    local_path = str(output_dir / instance_name)
    client.download_file(Bucket=bucket, Key=path, Filename=local_path)

def download_instances(
    client,
    instance_paths: Set[str],
    output_dir: Path,
    name_extractor: Callable[[str], str]
):
    for instance_path in tqdm(instance_paths):
        download_instance(client, instance_path, output_dir, name_extractor)

def select_batches(
    batches: Set[str],
    n: int,
    selection_criteria: str
):
    if selection_criteria == 'most_recent':
        batches = sorted(list(batches), key=int, reverse=True)
        batches = set(batches[:n])
    elif selection_criteria == 'random':
        batches = set(random.sample(batches, n))
    else:
        assert False
    return batches

def select_feasible_batches(
    max_nr_instances: int,
    max_fraction_infeasible: float,
    selection_criteria: str,
    result_paths: Path
):
    # select all feasible instances
    feasible = set()
    for name, path in instances(result_paths):
        try:
            instance = load_json(path)
        except json.JSONDecodeError:
            continue  # ignore empty files and other unparsable garbage
        if len(instance['orders']) > 0:
            feasible.add(name)

    feasible = {get_instance_batch(instance) for instance in feasible}

    # select subset of feasible instances
    nr_feasible = min(len(feasible), int(max_nr_instances * (1 - max_fraction_infeasible)))

    return select_batches(feasible, nr_feasible, selection_criteria)

def filter_unique(instances_path: Path):
    to_remove = set()
    for (name1, path1), (name2, path2) \
            in tqdm(combinations(instances(instances_path), 2)):
        if path1 in to_remove or path2 in to_remove:
            continue
        if filecmp.cmp(path1, path2):
            to_remove.add(path2)

    for path in to_remove:
        os.remove(path)

def get_batch_ids_from_file(filename):
    with open(filename, 'r') as f:
        return f.read().splitlines()

def filter_paths(paths, args):
    if args.from_date is not None:
        paths = filter_instances_by_date(paths, args.from_date)
    if args.with_batch_ids_from_file is not None:
        paths = filter_instances_by_batch(
            paths,
            get_batch_ids_from_file(args.with_batch_ids_from_file)
        )
    if args.with_batch_ids is not None:
        paths = filter_instances_by_batch(
            paths,
            args.with_batch_ids
        )
    return paths

def main(args):
    if args.max_fraction_infeasible != 1 and (
        args.with_batch_ids is not None
        or args.with_batch_ids_from_file is not None
    ):
        print("WARNING: '--max_fraction_infeasible' ignored when '--with_batch_ids' is given.")
        args.max_fraction_infeasible = 1

    instances_prefix = f'data/{args.network}/{args.solver}/instances/'
    results_prefix = f'data/{args.network}/{args.solver}/results/'

    client = boto3.client('s3')

    print("Downloading solution list.")
    result_paths = get_instance_list(client, prefix=results_prefix)
    result_paths = filter_paths(result_paths, args)

    batches = {get_instance_batch(result_path) for result_path in result_paths}

    result_dir = tempfile.TemporaryDirectory()

    print("Downloading solutions.")
    download_instances(client, result_paths, Path(result_dir.name), get_result_name)

    print("Selecting feasible instances.")
    feasible_batches = select_feasible_batches(
        args.max_nr_instances,
        args.max_fraction_infeasible,
        args.selection_criteria,
        Path(result_dir.name)
    )

    print("Downloading problem list.")
    instance_paths = get_instance_list(client, prefix=instances_prefix)
    instance_paths = filter_paths(instance_paths, args)

    print("Downloading feasible problems.")
    feasible_instance_paths = filter_instances_by_batch(instance_paths, feasible_batches)
    download_instances(
        client, feasible_instance_paths, Path(args.instances_path), get_instance_name
    )

    print("Selecting infeasible instances.")
    infeasible_batches = batches - feasible_batches
    if args.max_fraction_infeasible != 1:
        nr_infeasible = int(
            (len(feasible_batches) * args.max_fraction_infeasible)
            / (1 - args.max_fraction_infeasible)
        )
    else:
        nr_infeasible = len(infeasible_batches)

    infeasible_batches = select_batches(
        infeasible_batches,
        nr_infeasible,
        args.selection_criteria
    )

    print("Downloading infeasible problems.")
    infeasible_instance_paths = filter_instances_by_batch(
        instance_paths, infeasible_batches
    )
    download_instances(
        client, infeasible_instance_paths, Path(args.instances_path), get_instance_name
    )

    print("Removing duplicate problems.")
    filter_unique(Path(args.instances_path))

def setup_arg_parser(subparsers):
    parser = subparsers.add_parser('download', help="Downloads instances from S3 bucket.")

    parser.add_argument(
        'instances_path',
        type=Path,
        help='Directory where to copy instances to.'
    )
    parser.add_argument(
        '--from_date',
        type=str,
        help='Get instances newer than the given date (inclusive).'
    )
    parser.add_argument(
        '--max_nr_instances',
        type=int,
        default=100,
        help='Maximum number of instances to download.'
    )
    parser.add_argument(
        '--max_fraction_infeasible',
        type=float,
        default=1 / 3,
        help='Instances should have at most given fraction of infeasible instances.'
    )
    parser.add_argument(
        '--selection_criteria',
        type=str,
        choices=['most_recent', 'random'],
        default='random',
        help='How to select instances.'
    )
    parser.add_argument(
        '--with_batch_ids',
        type=str,
        nargs='+',
        help='Get instances with given batch ids (separated by spaces).'
    )
    parser.add_argument(
        '--with_batch_ids_from_file',
        type=str,
        help='Get instances with batch ids given in specified file (one per line).'
    )
    parser.add_argument(
        '--network',
        type=str,
        choices=['mainnet', 'mainnet-prod', 'rinkeby'],
        default='mainnet',
        help='Network where to pull instances from.'
    )
    parser.add_argument(
        '--solver',
        type=str,
        choices=['standard-solver', 'fallback-solver', 'open-solver'],
        default='standard-solver',
        help='Solver where to pull solution instances from.'
    )
    parser.set_defaults(exec_subcommand=main)


import importlib
class Device:
    def __init__(self, package_name):
        self.package_name = package_name
        self.dynamic_import()

    def dynamic_import(self):
        common_module = importlib.import_module('device_common')
        name_module_map = {name:common_module for name in common_module.__all__}
        if self.package_name != "device_common":
            special_module = importlib.import_module(self.package_name)
            name_module_map.update({name:special_module for name in special_module.__all__})
        for attr_name, module in name_module_map.items():
            class_module = importlib.import_module(name_module_map[attr_name].__name__+"."+attr_name)
            camel = "".join(map(str.capitalize, attr_name.split("_")))
            the_obj = getattr(class_module,camel)()
            setattr(self, attr_name, the_obj)

if __name__ == '__main__':
    from minitest import *

    with test(Device):
        d6000 = Device("device_6000")
        d6000.firewall.some().must_equal("device_common.firewall.Firewall")
        d6000.system.some().must_equal("device_6000.system.System")

        d7000 = Device("device_7000")
        d7000.firewall.some().must_equal("device_7000.firewall.Firewall")
        d7000.load_balance.some().must_equal("device_7000.load_balance.LoadBalance")
        d7000.system.some().must_equal("device_7000.system.System")
        d7000.system.common().must_equal("common")
        d7000.system.special().must_equal("special")


from meta_policy_search.baselines.linear_baseline import LinearFeatureBaseline, MetaNNBaseline, LinearTimeBaseline
from meta_policy_search.envs.mujoco_envs.humanoid_rand_direc import HumanoidRandDirecEnv
from meta_policy_search.envs.mujoco_envs.humanoid_rand_direc_2d import HumanoidRandDirec2DEnv
from meta_policy_search.envs.mujoco_envs.ant_rand_direc import AntRandDirecEnv
from meta_policy_search.envs.mujoco_envs.ant_rand_direc_2d import AntRandDirec2DEnv
from meta_policy_search.envs.mujoco_envs.ant_rand_goal import AntRandGoalEnv
from meta_policy_search.envs.mujoco_envs.half_cheetah_rand_direc import HalfCheetahRandDirecEnv
from meta_policy_search.envs.mujoco_envs.half_cheetah_rand_vel import HalfCheetahRandVelEnv
from meta_policy_search.envs.mujoco_envs.humanoid_rand_direc import HumanoidRandDirecEnv
from meta_policy_search.envs.mujoco_envs.humanoid_rand_direc_2d import HumanoidRandDirec2DEnv
from meta_policy_search.envs.mujoco_envs.walker2d_rand_direc import Walker2DRandDirecEnv
from meta_policy_search.envs.mujoco_envs.walker2d_rand_vel import Walker2DRandVelEnv
from meta_policy_search.envs.mujoco_envs.swimmer_rand_vel import SwimmerRandVelEnv
from meta_policy_search.envs.normalized_env import normalize
from meta_policy_search.meta_algos.tmaml import TMAML
from meta_policy_search.meta_trainer import Trainer
from meta_policy_search.samplers.meta_sampler import MetaSampler
from meta_policy_search.samplers.tmaml_sample_processor import TMAMLMetaSampleProcessor
from meta_policy_search.policies.meta_gaussian_mlp_policy import MetaGaussianMLPPolicy
from meta_policy_search.utils import logger
from meta_policy_search.utils.utils import set_seed, ClassEncoder

import numpy as np
import os
import json
import argparse
import time
import os

def main(config):
    set_seed(config['seed'])

    reward_baseline = LinearTimeBaseline() # the usual baseline
    return_baseline = LinearFeatureBaseline() # the additional baseline for DICE

    env = globals()[config['env']]() # instantiate env
    env = normalize(env) # apply normalize wrapper to env

    meta_baseline = MetaNNBaseline(input_size=env.observation_space.shape[0]) # the meta baseline

    policy = MetaGaussianMLPPolicy(
            name="meta-policy",
            obs_dim=np.prod(env.observation_space.shape),
            action_dim=np.prod(env.action_space.shape),
            meta_batch_size=config['meta_batch_size'],
            hidden_sizes=config['hidden_sizes'],
    )

    sampler = MetaSampler(
        env=env,
        policy=policy,
        rollouts_per_meta_task=config['rollouts_per_meta_task'],
        meta_batch_size=config['meta_batch_size'],
        max_path_length=config['max_path_length'],
        parallel=config['parallel'],
    )

    sample_processor = TMAMLMetaSampleProcessor(
        baseline=reward_baseline,
        max_path_length=config['max_path_length'],
        discount=config['discount'],
        normalize_adv=config['normalize_adv'],
        positive_adv=config['positive_adv'],
        return_baseline=return_baseline,
        metabaseline=meta_baseline,
    )

    algo = TMAML(
        policy=policy,
        max_path_length=config['max_path_length'],
        meta_batch_size=config['meta_batch_size'],
        num_inner_grad_steps=config['num_inner_grad_steps'],
        inner_lr=config['inner_lr'],
        learning_rate=config['learning_rate']
    )

    trainer = Trainer(
        algo=algo,
        policy=policy,
        env=env,
        sampler=sampler,
        sample_processor=sample_processor,
        n_itr=config['n_itr'],
        num_inner_grad_steps=config['num_inner_grad_steps'],
    )

    trainer.train()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Taming MAML')
    parser.add_argument('--dump_path', type=str, default=os.path.join(os.path.realpath(os.path.dirname(__file__)), 'logs'))
    parser.add_argument('--seed', type=int, help='random seed', default=0)
    parser.add_argument('--env', type=str,
                        help='env name', default='HalfCheetahRandDirecEnv')
    parser.add_argument('--inner_type', type=str, help='type of inner loss function', default='log_likelihood')
    parser.add_argument('--n_iter', type=int,
                        help='number of iterations', default=5001)

    args = parser.parse_args()

    os.makedirs(args.dump_path, exist_ok=True)

    config = {
        'seed': 1,

        'baseline': 'LinearFeatureBaseline',

        'env': 'HalfCheetahRandDirecEnv',

        # sampler config
        'rollouts_per_meta_task': 20,
        'max_path_length': 100,
        'parallel': True,

        # sample processor config
        'discount': 0.99,
        'gae_lambda': 1,
        'normalize_adv': True,
        'positive_adv': False,

        # policy config
        'hidden_sizes': (64, 64),
        'learn_std': True,  # whether to learn the standard deviation of the gaussian policy

        'inner_lr': 0.1,  # adaptation step size
        'learning_rate': 1e-3,  # meta-policy gradient step size
        'step_size': 0.01,  # size of the TRPO trust-region
        'n_itr': 5001,  # number of overall training iterations
        'meta_batch_size': 40,  # number of sampled meta-tasks per iterations
        'num_inner_grad_steps': 1,  # number of inner / adaptation gradient steps
        'inner_type': 'log_likelihood',  # type of inner loss function used

    }

    config.update(vars(args))

    # configure logger
    logger.configure(dir=args.dump_path, format_strs=['stdout', 'log', 'csv', 'tensorboard'],
                     snapshot_mode='last_gap')

    # dump run configuration before starting training
    json.dump(config, open(args.dump_path +
                           '/params.json', 'w'), cls=ClassEncoder)

    # start the actual algorithm
    main(config)


import random
from statistics import median

def fill(ax,nx, x,y):       # fill the list a with n random integers in the range x --- y inclusive BUT NO duplicates permitted THIS code NEEDS to be fixed
    j = 0
    while (j < nx):  # populate the array using the while loop
        a = random.randint(x, y)  # Return a random integer .
        if a not in ax:
          ax.append(a)
          j = j + 1
    return ax
def display(ax):
    print("2. Display the list as follows: ", ax)  # DONE

def sort_list_asc(ax):
    sorted_list = sorted(ax, reverse=True)
    return sorted_list

def min_integer(ax):
    # return the smallest value in the list
    smallest = min(ax)
    return smallest

def max_integer(ax):
    # return the largest value in the list
    mx = max(ax)  # place holder only
    return mx

def div3(ax):
    count_three = 0
    for i in ax:
        if i % 3 == 0:
            count_three = count_three + 1
    return count_three
    # v = -5
    # return v

def multi5(ax):
    count_five = 0
    for i in ax:
        if i % 5 == 0:
            count_five = count_five + 1
    return count_five
    # m = 11
    # return m

def avg_of_list(ax):
    # return the average of the integers in the list
    #must call the sum() function here and divide by the size of the array
    size_of_array = sum(ax)
    average_list = size_of_array/len(ax)
    return average_list

def maximum(ax):
    # return the largest value in the list
    mx = ax[0]  # place holder only
    return mx

def minimum(ax):
    # return the smallest value in the list
    mn = ax[-1]
    return mn

def median_number(ax):
    med = median(ax)
    return med

def digit1(ax):
    list_ones = [11,21,31]
    count_end_one = 0
    for i in ax:
        if i in list_ones:
            count_end_one = count_end_one + 1
    return count_end_one

# ==========================================
# NOW we call the functions
#Suggested template i.e all calculations should be in the functions; simply call the functions
array = []
n = 21  # array size
x = 10
y = 40
array = fill(array,n, x, y)
display(array)
mx = max_integer(array)
mn = min_integer(array)
d = div3(array)
m = multi5(array)
med = median_number(array)
n1 = digit1(array)
array = sort_list_asc(array)
ave = avg_of_list(array)
print("3. The sorted list in descending order is ", array)
print( "4. The maximum integer in the list  is  : ", mx )
print( "5. The minimum integer in the list  is :", mn)
print( "6. The number of integers that are evenly divisible by 3 (no remainder)is : ", d)
print( "7. The number of integers in the list that are multiples of 5 (10, 15, 20,25, etc ) is : ", m)
print( "8. The median of the list is : ", med)
print( "9. the number of integers that end with a 1 is : ", n1)
print("10. The average of the integers is ", round(ave), ".( Please note the number is rounded.)")


import scrapy
from scrapy.http import FormRequest

from locations.geo import point_locations
from locations.structured_data_spider import StructuredDataSpider

class PetsAtHomeGBSpider(StructuredDataSpider):
    name = "petsathome_gb"
    item_attributes = {"brand": "Pets at Home", "brand_wikidata": "Q7179258"}
    wanted_types = ["PetStore"]
    download_delay = 0.2

    def start_requests(self):
        for lat, lon in point_locations("eu_centroids_20km_radius_country.csv", "UK"):
            yield FormRequest(
                url="https://community.petsathome.com/models/utils.cfc",
                formdata={
                    "method": "functionhandler",
                    "returnFormat": "json",
                    "event": "webproperty.storelocator",
                    "lat": str(lat),
                    "lng": str(lon),
                    "radius": "50",
                    "companyID": "any",
                    "active": "true",
                },
                headers={"X-Requested-With": "XMLHttpRequest"},
            )

    def parse(self, response):
        for node in response.json()["data"]:
            if "groom-room" not in node["slug"]:
                website_url = "https://community.petsathome.com" + node["slug"]
                yield scrapy.Request(website_url, callback=self.parse_sd)


"""
    Author: pablopg.
    Step-by-step:
        1. Run pip install boto3
        2. Run pip install --upgrade --user awscli (http://docs.aws.amazon.com/cli/latest/userguide/installing.html)
        3. Add awscli to the path (http://docs.aws.amazon.com/cli/latest/userguide/awscli-install-windows.html#awscli-install-windows-path)
        4. Create an user in the AWS CLI.
        5. Configure local AWS cli via the command: aws configure --profile adminuser
        6. Run and listen!
"""
"""Getting Started Example for Python 2.7+/3.3+"""
from boto3 import Session
from botocore.exceptions import BotoCoreError, ClientError
from contextlib import closing
import os
import sys
import subprocess
from tempfile import gettempdir
import audio

# Create a client using the credentials and region defined in the [adminuser]
# section of the AWS credentials file (~/.aws/credentials).
session = Session(profile_name="default", region_name='eu-west-1')
polly = session.client("polly")
id_n = 0

def polly_play(input_str):
    global id_n
    id_n = id_n + 1
    try:
        # Request speech synthesis
        response = polly.synthesize_speech(Text=input_str, OutputFormat="mp3",
                                           VoiceId="Kendra")
    except (BotoCoreError, ClientError) as error:
        # The service returned an error, exit gracefully
        print(error)
        sys.exit(-1)

    # Access the audio stream from the response
    if "AudioStream" in response:
        # Note: Closing the stream is important as the service throttles on the
        # number of parallel connections. Here we are using contextlib.closing to
        # ensure the close method of the stream object will be called automatically
        # at the end of the with statement's scope.
        with closing(response["AudioStream"]) as stream:
            output = os.path.join(gettempdir(), "speech"+str(id_n)+".mp3")

            try:
                # Open a file for writing the output as a binary stream
                with open(output, "wb") as file:
                    file.write(stream.read())
                    file.close()
            except IOError as error:
                # Could not write to file, exit gracefully
                print(error)
                sys.exit(-1)

    else:
        # The response didn't contain audio data, exit gracefully
        print("Could not stream audio")
        sys.exit(-1)

    # Play the audio using python

    audio.play_mp3(output, 1.0)


# -*- coding: utf-8 -*-
from datetime import timedelta, datetime
from openerp import models, fields, api, exceptions, _
from openerp.exceptions import ValidationError
from dateutil.relativedelta import relativedelta
from openerp.osv import osv
import openerp.addons.decimal_precision as dp
import time
import tools

class academic_year(models.Model):
    _name = 'oschool.academic_year'
    _rec_name = 'code'

    name = fields.Char(string="Designation", required=True)
    code = fields.Char(string="Code", required=True)
    date_start = fields.Date(string="Start Date", required=True)
    date_stop = fields.Date(string="End Date", required=True)
    period_ids = fields.Many2many('account.period', string="Period")
    company_id = fields.Many2one('res.company', 'Company')
    state = fields.Selection([
                                 ('new', "New"),
                                 ('current', "Current"),
                                 ('closed', "Closed"),
                             ], default='new')

    def _generate_periods(self,cr,uid, date_start,date_stop, context=None):
        current_user = self.pool.get('res.users').browse(cr, uid, uid, context=context)
        periods = self.pool.get('account.period').search(cr,uid,[('date_start','>=',date_start),('date_start','<=',date_stop), ('code','not ilike','00'),('company_id','=',current_user.company_id.id)])
        tmp =[]
        for period in periods:
            tmp.append(period)
        return [[6, False, tmp]]

    def create(self, cr, uid, vals, context=None):

        date_start = datetime.strptime(vals['date_start'], '%Y-%m-%d')
        date_stop = datetime.strptime(vals['date_stop'], '%Y-%m-%d')
        dat_start = "%s-%s-01" % (date_start.year,date_start.month)
        date_stop = "%s-%s-01" % (date_stop.year,date_stop.month)
        vals['period_ids'] = self._generate_periods(cr,uid,dat_start,date_stop)
        current_user = self.pool.get('res.users').browse(cr, uid, uid, context=context)
        current = self.search(cr, uid, [('state', '=', 'new'),('company_id','=', current_user.company_id.id)], context=context)
        if len(current) > 0:
            raise ValidationError("School year with the new state already exist")
        new_id = super(academic_year, self).create(cr, uid, vals, context=context)
        return new_id

    def write(self, cr, uid, ids, values, context = None):

        if values.has_key('date_start'):
            date_start = datetime.strptime(values['date_start'], '%Y-%m-%d')
        else:
            for fy in self.browse(cr, uid, ids, context=context):
                date_start = datetime.strptime(fy.date_start, '%Y-%m-%d')
        if values.has_key('date_stop'):
            date_stop = datetime.strptime(values['date_stop'], '%Y-%m-%d')
        else:
            for fy in self.browse(cr, uid, ids, context=context):
                date_stop = datetime.strptime(fy.date_stop, '%Y-%m-%d')

        dat_start = "%s-%s-01" % (date_start.year,date_start.month)
        date_stop = "%s-%s-01" % (date_stop.year,date_stop.month)
        values['period_ids'] =  self._generate_periods(cr,uid,dat_start,date_stop)
        res = super(academic_year, self).write(cr, uid, ids, values, context = context)
        return res

    def action_new(self, cr, uid, ids, context=None):
        current_user = self.pool.get('res.users').browse(cr, uid, uid, context=context)
        current = self.search(cr, uid, [('state', '=', 'new'),('company_id','=',current_user.company_id.id)])
        if len(current) > 0:
            raise ValidationError("School year with the new state already exist")
        return self.write(cr, uid, ids, {'state': 'new'})

    def action_current(self, cr, uid, ids, context=None):
        return self.write(cr, uid, ids, {'state': 'current'})

    def action_closed(self, cr, uid, ids, context=None):
        return self.write(cr, uid, ids, {'state': 'closed'})

    _sql_constraints = [
        ('code_unique',
         'UNIQUE(code,company_id)',
         "The code must be unique"),
        ('name_unique',
         'UNIQUE(name,company_id)',
         "The name must be unique"),
        ('check_date',
         'CHECK(date_start < date_stop)',
         "Start Date Must be Below End Date"),
    ]

    _defaults = {
        'company_id': tools.get_default_company,
    }

class groups(models.Model):
    _name = 'oschool.groups'

    name = fields.Char(string="Group designation", required=True)
    code = fields.Char(string="Code", required=True)
    number_of_places = fields.Integer(string="Number of places", required=True)
    class_ids = fields.One2many('oschool.classes', 'code', 'Classes')
    seq = fields.Integer(string="Sequence",required=True, help="This field determine the order of the groups, is required for transport report")
    company_id = fields.Many2one('res.company', 'Company')

    _defaults = {
        'company_id': tools.get_default_company,
    }

    _sql_constraints = [
        ('code_unique',
         'UNIQUE(code,company_id)',
         "Group code must be unique"),
        ('name_unique',
         'UNIQUE(name,company_id)',
         "The name must be unique"),
        ('seq_unique',
         'UNIQUE(seq)',
         "Group sequence must be unique")
    ]

class classes(models.Model):
    _name = 'oschool.classes'

    name = fields.Char(string="Classroom designation", required=True)
    code = fields.Char(string="Code", required=True)
    group = fields.Many2one('oschool.groups', 'Groupe', ondelete='cascade')
    company_id = fields.Many2one('res.company', 'Company')

    _defaults = {
        'company_id': tools.get_default_company,
    }
    _sql_constraints = [
        ('code_unique',
         'UNIQUE(code,company_id)',
         "Le code de la classe doit etre unique"),
        ('name_unique',
         'UNIQUE(name,company_id)',
         "The name must be unique")
    ]

class pos_category(models.Model):
    _inherit = 'pos.category'

    price = fields.Float(string="Price")
    place_number = fields.Integer(string="place number")
    academic_year = fields.Many2one('oschool.academic_year', ondelete='cascade', string="Academic year")
    start_date_registration = fields.Date(string='Start Registration Date')
    stop_date_registration = fields.Date(string='End Registration Date')
    date_start = fields.Date(string='Start Date')
    date_stop = fields.Date(string='End Date')
    services_ids = fields.One2many('product.template', 'pos_categ_id', ondelete='cascade', string='Months')
    product_category = fields.Many2one('product.category', string="Category Type")
    cash = fields.Boolean(default=False, string="Cash")
    school_product_type = fields.Selection(
        [('service', 'Service'), ('registration', 'Registration'), ('club', 'Club'), ('event', 'Event'),
         ('study', 'Study'), ('extra', 'Extra'), ('transport', 'Transport')])
    groups = fields.Many2one('oschool.groups', 'Groupe', ondelete='cascade')
    excluded_services_ids = fields.Many2many('pos.category', 'service_category_rel', 'category_id', 'service_id', string='Excluded services')
    company_id = fields.Many2one('res.company', 'Company')

    @api.one
    @api.onchange('academic_year')
    def _set_start_stop_date(self):
        self.date_start = self.academic_year.date_start
        self.date_stop = self.academic_year.date_stop

    _sql_constraints = [
        ('check_date',
         'CHECK(date_start <= date_stop)',
         "Start Date Must be Below End Date"),
        ('check_registration_date',
         'CHECK(start_date_registration < stop_date_registration)',
         "Start Date Registration Must be Below End Date registration"),
        ('check_registration_date_stop_date',
         'CHECK(stop_date_registration <= date_stop)',
         "Stop Date Registration Must be Below End Date"),
        ('name_unique',
         'UNIQUE(name,academic_year)',
         "This record already exists for this academic year"),
    ]

    _defaults = {
        'company_id': tools.get_default_company,
    }

    def generate_months_clubs(self, cr, uid, ids, context=None, interval=1):
        company_id = self.pool.get('res.users')._get_company(cr, uid, context=context)
        property_account_income_extra = self.pool.get('account.account').search(cr, uid, [('code', '=', '7050005'),('company_id','=',company_id)])
        if not property_account_income_extra:
            raise ValidationError("There is not account with code 7050005")
        period_obj = self.pool.get('product.template')
        for fy in self.browse(cr, uid, ids, context=context):
            ds = datetime.strptime(fy.date_start, '%Y-%m-%d')
            while ds.strftime('%Y-%m-%d') < fy.date_stop:
                de = ds + relativedelta(months=interval, days=-1)
                if de.strftime('%Y-%m-%d') > fy.date_stop:
                    de = datetime.strptime(fy.date_stop, '%Y-%m-%d')

                period_obj.create(cr, uid, {
                    'name': "%s %s" % (fy.name, ds.strftime('%m/%Y')),
                    'pos_categ_id': fy.id,
                    'type': 'service',
                    'available_in_pos': True,
                    'property_account_income': property_account_income_extra[0],
                    'list_price': fy.price,
                    'subscription_month': ds.strftime('%m/%Y'),
                    'school_product_type': fy.school_product_type,
                    'academic_year': fy.academic_year.id,
                    'categ_id': fy.product_category.id,
                    'cash': fy.cash,
                    'company_id': fy.company_id.id,
                })
                ds = ds + relativedelta(months=interval)
        return True

    def generate_inscriptions(self, cr, uid, ids, context=None, interval=1):
        company_id = self.pool.get('res.users')._get_company(cr, uid, context=context)
        property_account_income_inscription = self.pool.get('account.account').search(cr, uid, [('code', '=', '7050004'),('company_id','=',company_id)])
        if not property_account_income_inscription:
            raise ValidationError("There is not account with code 7050004")
        period_obj = self.pool.get('product.template')
        sequence = 1
        for fy in self.browse(cr, uid, ids, context=context):
            if (fy.academic_year.state) == "current" or (fy.academic_year.state) == "closed":
                error = "You cannot generate inscriptions year that have '%s' status" % _(fy.academic_year.state)
                raise ValidationError(error)
            period_obj.create(cr, uid, {
                'name': "%s %s" % (_("Inscription"), fy.academic_year.code),
                'pos_categ_id': fy.id,
                'type': 'service',
                'available_in_pos': True,
                'property_account_income': property_account_income_inscription[0],
                'school_product_type': fy.school_product_type,
                'academic_year': fy.academic_year.id,
                'categ_id': fy.product_category.id,
                'inscription_sequence': sequence,
            })
            sequence = sequence + 1
            period_obj.create(cr, uid, {
                'name': "%s %s" % (_("Rénscription"), fy.academic_year.code),
                'pos_categ_id': fy.id,
                'type': 'service',
                'available_in_pos': True,
                'property_account_income': property_account_income_inscription[0],
                'school_product_type': fy.school_product_type,
                'academic_year': fy.academic_year.id,
                'categ_id': fy.product_category.id,
                'inscription_sequence': sequence,
            })

    def generate_months_study(self, cr, uid, ids, context=None, interval=1):
        company_id = self.pool.get('res.users')._get_company(cr, uid, context=context)
        property_account_income = self.pool.get('account.account').search(cr, uid, [('code', '=', '7050001'),('company_id','=',company_id)])
        if not property_account_income:
            raise ValidationError("There is not account with code 7050001")
        period_obj = self.pool.get('product.template')
        for fy in self.browse(cr, uid, ids, context=context):
            ds = datetime.strptime(fy.academic_year.date_start, '%Y-%m-%d')
            while ds.strftime('%Y-%m-%d') < fy.academic_year.date_stop:
                de = ds + relativedelta(months=interval, days=-1)
                if de.strftime('%Y-%m-%d') > fy.academic_year.date_stop:
                    de = datetime.strptime(fy.academic_year.date_stop, '%Y-%m-%d')

                period_obj.create(cr, uid, {
                    'name': "%s %s" % (fy.name, ds.strftime('%m/%Y')),
                    'pos_categ_id': fy.id,
                    'type': 'service',
                    'available_in_pos': True,
                    'list_price': fy.price,
                    'property_account_income': property_account_income[0],
                    'subscription_month': ds.strftime('%m/%Y'),
                    'school_product_type': fy.school_product_type,
                    'academic_year': fy.academic_year.id,
                    'cash': fy.cash,
                    'categ_id': fy.product_category.id,

                })
                ds = ds + relativedelta(months=interval)

    def generate_months_service(self, cr, uid, ids, context=None, interval=1):
        company_id = self.pool.get('res.users')._get_company(cr, uid, context=context)
        property_account_income = self.pool.get('account.account').search(cr, uid, [('code', '=', '7050002'),('company_id','=',company_id)])
        if not property_account_income:
            raise ValidationError("There is not account with code 7050002")
        period_obj = self.pool.get('product.template')
        for fy in self.browse(cr, uid, ids, context=context):
            ds = datetime.strptime(fy.date_start, '%Y-%m-%d')
            while ds.strftime('%Y-%m-%d') < fy.date_stop:
                de = ds + relativedelta(months=interval, days=-1)
                if de.strftime('%Y-%m-%d') > fy.date_stop:
                    de = datetime.strptime(fy.date_stop, '%Y-%m-%d')

                period_obj.create(cr, uid, {
                    'name': "%s %s" % (fy.name, ds.strftime('%m/%Y')),
                    'pos_categ_id': fy.id,
                    'type': 'service',
                    'available_in_pos': True,
                    'list_price': fy.price,
                    'property_account_income': property_account_income[0],
                    'subscription_month': ds.strftime('%m/%Y'),
                    'school_product_type': fy.school_product_type,
                    'categ_id': fy.product_category.id,
                    'academic_year': fy.academic_year.id,
                    'cash': fy.cash,
                })
                ds = ds + relativedelta(months=interval)

    def unlink(self, cr, uid, ids, context=None):
        product_obj = self.pool.get('product.product')
        product_tmpl_obj = self.pool.get('product.template')
        pos_order_line_obj = self.pool.get('pos.order.line')
        for id in ids:
            category = self.browse(cr, uid, id)
            product_tmpl_ids = product_tmpl_obj.search(cr, uid,
                                    [('pos_categ_id','=',id),
                                     ('school_product_type','=',category.school_product_type)])

            for product_tmpl_id in product_tmpl_ids:
                product_ids = product_obj.search(cr, uid,[('product_tmpl_id','=',product_tmpl_id)])
                for product_id in product_ids:
                    pos_order_line_ids = pos_order_line_obj.search(cr, uid,[('product_id','=',product_id)])
                    if pos_order_line_ids:
                        #product = product_obj.browse(cr, uid, product_id)
                        msg = "One or more command was created for the category {0}.".format(category.name)
                        raise exceptions.Warning(msg)
                        return False
            product_tmpl_obj.unlink(cr, uid, product_tmpl_ids, context)
            super(pos_category, self).unlink(cr, uid, id, context)
        return True

class product_template(models.Model):
    _inherit = 'product.template'

    subscription_month = fields.Char(string='Subscription Month')
    cash = fields.Boolean(default=False, string="Cash")
    school_product_type = fields.Selection(
        [('service', 'Service'), ('registration', 'Registration'), ('club', 'Club'), ('event', 'Event'),
         ('study', 'Study'), ('extra', 'Extra'), ('transport', 'Transport')])
    inscription_sequence = fields.Integer(string="sequence")

class hr_employee(models.Model):
    _name = 'hr.employee'
    _inherit = 'hr.employee'
    _description = 'Driver Information'

    licence_no = fields.Char(string='Licence No', size=50)
    company_id = fields.Many2one('res.company', 'Company')

    _defaults = {
        'company_id': tools.get_default_company,
    }

class parents_categories(models.Model):
    _inherit = 'product.pricelist.item'

    _defaults = {
        'base': 1,
    }

class oschool_account_journal(models.Model):
    _inherit = 'account.journal'

    cash = fields.Boolean(default=False, string="Cash")
    registration = fields.Boolean(default=False, string="Registration")
    extra = fields.Boolean(default=False, string="Extra")
    is_check = fields.Boolean(default=False, string="Is check")

class pos_session_opening(osv.osv_memory):
    _inherit = 'pos.session.opening'

    company_id = fields.Many2one('res.company', 'Company')

    _defaults = {
        'company_id': tools.get_default_company,
    }
    def oschool_open_session_cb(self, cr, uid, ids, context=None):
        assert len(ids) == 1, "you can open only one session at a time"
        proxy = self.pool.get('pos.session')
        wizard = self.browse(cr, uid, ids[0], context=context)
        if not wizard.pos_session_id:
            values = {
                'user_id': uid,
                'config_id': wizard.pos_config_id.id,
            }
            session_id = proxy.create(cr, uid, values, context=context)
            s = proxy.browse(cr, uid, session_id, context=context)
            if s.state == 'opened':
                return self.oschool_open_ui(cr, uid, ids, context=context)
            return self._open_session(session_id)
        return self._open_session(wizard.pos_session_id.id)

    def oschool_open_ui(self, cr, uid, ids, context=None):
        data = self.browse(cr, uid, ids[0], context=context)
        context = dict(context or {})
        context['active_id'] = data.pos_session_id.id
        obj = self.pool.get('ir.model.data').get_object_reference(cr, uid, 'oschool', 'main_oschool_menu')[1]
        return {
            'type': 'ir.actions.client',
            'name': 'Point of Sale Menu',
            'tag': 'reload',
            'params': {'menu_id': obj},
        }

    def oschool_open_existing_session_cb_close(self, cr, uid, ids, context=None):
        wizard = self.browse(cr, uid, ids[0], context=context)
        wizard.pos_session_id.signal_workflow('cashbox_control')
        return self.oschool_open_session_cb(cr, uid, ids, context)

    def oschool_open_session_cb(self, cr, uid, ids, context=None):

        view_id = self.pool.get('ir.model.data').get_object_reference(cr, uid, 'oschool', 'oschool_pos_session_form')

        assert len(ids) == 1, "you can open only one session at a time"
        proxy = self.pool.get('pos.session')
        wizard = self.browse(cr, uid, ids[0], context=context)
        if not wizard.pos_session_id:
            values = {
                'user_id': uid,
                'config_id': wizard.pos_config_id.id,
            }
            session_id = proxy.create(cr, uid, values, context=context)
            s = proxy.browse(cr, uid, session_id, context=context)
            if s.state == 'opened':
                return self.oschool_open_ui(cr, uid, ids, context=context)
            return self._oschool_open_session(session_id, view_id, context)
        return self._oschool_open_session(wizard.pos_session_id.id, view_id,context)

    def _oschool_open_session(self, session_id, view_id, context):

        return {
            'name': _('Session'),
            'view_type': 'form',
            'view_mode': 'form,tree',
            'res_model': 'pos.session',
            'res_id': session_id,
             # 'view_id': [view_id],
            'view_id': False,
            'type': 'ir.actions.act_window',
        }

class pos_session(osv.osv):

    _inherit = 'pos.session'

    company_id = fields.Many2one('res.company', 'Company')

    _defaults = {
        'company_id': tools.get_default_company,
    }

    def wkf_action_close(self, cr, uid, ids, context=None):
        # Close CashBox
        for record in self.browse(cr, uid, ids, context=context):
            for st in record.statement_ids:
                if abs(st.difference) > st.journal_id.amount_authorized_diff:
                    # The pos manager can close statements with maximums.
                    if not self.pool.get('ir.model.access').check_groups(cr, uid, "point_of_sale.group_pos_manager"):
                        raise osv.except_osv(_('Error!'),
                                             _(
                                                 "Your ending balance is too different from the theoretical cash closing (%.2f), the maximum allowed is: %.2f. You can contact your manager to force it.") % (
                                             st.difference, st.journal_id.amount_authorized_diff))
                if (st.journal_id.type not in ['bank', 'cash']):
                    raise osv.except_osv(_('Error!'),
                                         _("The type of the journal for your payment method should be bank or cash "))
                getattr(st, 'button_confirm_%s' % st.journal_id.type)(context=context)
        self._confirm_orders(cr, uid, ids, context=context)
        self.write(cr, uid, ids, {'state': 'closed'}, context=context)

        obj = self.pool.get('ir.model.data').get_object_reference(cr, uid, 'oschool', 'main_oschool_menu')[1]
        return {
            'type': 'ir.actions.client',
            'name': 'Point of Sale Menu',
            'tag': 'reload',
            'params': {'menu_id': obj},
        }

    def open_frontend_cb(self, cr, uid, ids, context=None):
        if not context:
            context = {}
        if not ids:
            return {}
        for session in self.browse(cr, uid, ids, context=context):
            if session.user_id.id != uid:
                raise osv.except_osv(
                    _('Error!'),
                    _(
                        "You cannot use the session of another users. This session is owned by %s. Please first close this one to use this point of sale." % session.user_id.name))
        context.update({'active_id': ids[0]})
        obj = self.pool.get('ir.model.data').get_object_reference(cr, uid, 'oschool', 'main_oschool_menu')[1]
        return {
            'type': 'ir.actions.client',
            'name': 'Point of Sale Menu',
            'tag': 'reload',
            'params': {'menu_id': obj},
        }

class product_template(models.Model):
    _inherit = 'product.template'

    academic_year = fields.Many2one('oschool.academic_year', ondelete='cascade', string="Academic year")

    _sql_constraints = [
        ('unique_oschool_product', 'unique(academic_year,company_id, school_product_type, subscription_month, pos_categ_id)',
         'The record must be unique by year'),
        ('unique_oschool_inscription_by_year',
         'unique(academic_year, company_id,school_product_type, inscription_sequence, categ_id)',
         'The registration must be unique by year')
    ]


# -*- coding: utf-8 -*-
#
# Copyright (C) 2020 CERN.
# Copyright (C) 2020 Northwestern University.
#
# Flask-Resources is free software; you can redistribute it and/or modify
# it under the terms of the MIT License; see LICENSE file for more details.

"""Test deserialization."""

import marshmallow as ma
from flask import Flask, request
from speaklater import make_lazy_string

from flask_resources.serializers import JSONSerializer, MarshmallowJSONSerializer

def _(s):
    return make_lazy_string(lambda: s)

def test_lazy_strings_are_serialized():
    serializer = JSONSerializer()
    lazy = {"key": _("Lazy")}
    list_lazy = [{"key": _("Lazy1")}, {"key": _("Lazy2")}]

    assert '{"key": "Lazy"}' == serializer.serialize_object(lazy)
    assert '[{"key": "Lazy1"}, {"key": "Lazy2"}]' == serializer.serialize_object_list(
        list_lazy
    )

def test_prettyprint():
    app = Flask("test")
    with app.test_request_context("/?prettyprint=1"):
        serializer = JSONSerializer()
        assert '{\n  "key": "1"\n}' == serializer.serialize_object({"key": "1"})

def test_marhsmallow_serializer():
    class TestSchema(ma.Schema):
        title = ma.fields.Str(data_key="test")

    s = MarshmallowJSONSerializer(item_schema=TestSchema, list_schema=TestSchema)
    assert s.serialize_object({"title": "a"}) == '{"test": "a"}'
    assert s.serialize_object_list({"title": "a"}) == '{"test": "a"}'