"""
Example Thetis configuration with traitlets
"""

from traitlets.config.configurable import Configurable
from traitlets import *

from thetis import FiredrakeConstant as Constant
from thetis import FiredrakeFunction as Function


class PositiveInteger(Integer):
    def info(self):
        return u'a positive integer'

    def validate(self, obj, proposal):
        super(PositiveInteger, self).validate(obj, proposal)
        assert proposal > 0, self.error(obj, proposal)
        return proposal


class PositiveFloat(Float):
    def info(self):
        return u'a positive float'

    def validate(self, obj, proposal):
        super(PositiveFloat, self).validate(obj, proposal)
        assert proposal > 0.0, self.error(obj, proposal)
        return proposal


class BoundedInteger(Integer):
    def __init__(self, default_value=Undefined, bounds=None, **kwargs):
        super(BoundedInteger, self).__init__(default_value, **kwargs)
        self.minval = bounds[0]
        self.maxval = bounds[1]

    def info(self):
        return u'an integer between {:} and {:}'.format(self.minval, self.maxval)

    def validate(self, obj, proposal):
        super(BoundedInteger, self).validate(obj, proposal)
        assert proposal >= self.minval, self.error(obj, proposal)
        assert proposal <= self.maxval, self.error(obj, proposal)
        return proposal


class BoundedFloat(Float):
    def __init__(self, default_value=Undefined, bounds=None, **kwargs):
        self.minval = bounds[0]
        self.maxval = bounds[1]
        super(BoundedFloat, self).__init__(default_value, **kwargs)

    def info(self):
        return u'a float between {:} and {:}'.format(self.minval, self.maxval)

    def validate(self, obj, proposal):
        super(BoundedFloat, self).validate(obj, proposal)
        assert proposal >= self.minval, self.error(obj, proposal)
        assert proposal <= self.maxval, self.error(obj, proposal)
        return proposal


class FiredrakeConstant(TraitType):
    default_value = None
    info_text = 'a Firedrake Constant'

    def validate(self, obj, value):
        if isinstance(value, Constant):
            return value
        self.error(obj, value)

    def default_value_repr(self):
        return 'Constant({:})'.format(self.default_value.dat.data[0])


class FiredrakeCoefficient(TraitType):
    default_value = None
    info_text = 'a Firedrake Constant or Function'

    def validate(self, obj, value):
        if isinstance(value, (Constant, Function)):
            return value
        self.error(obj, value)

    def default_value_repr(self):
        if isinstance(self.default_value, Constant):
            return 'Constant({:})'.format(self.default_value.dat.data[0])
        return 'Function'


class FrozenHasTraits(HasTraits):
    """
    A HasTraits class that only allows adding new attributes in the class
    definition or when  self._isfrozen is False.
    """
    _isfrozen = False
    def __init__(self, *args, **kwargs):
        super(FrozenHasTraits, self).__init__(*args, **kwargs)
        self._isfrozen = True

    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(FrozenHasTraits, self).__setattr__(key, value)


class FrozenConfigurable(Configurable):
    """
    A Configurable class that only allows adding new attributes in the class
    definition or when  self._isfrozen is False.
    """
    _isfrozen = False
    def __init__(self, *args, **kwargs):
        super(FrozenConfigurable, self).__init__(*args, **kwargs)
        self._isfrozen = True

    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(FrozenConfigurable, self).__setattr__(key, value)


class TimeStepperOptions(FrozenHasTraits):
    """Base class for all time stepper options"""


class CrankNicolsonOptions(TimeStepperOptions):
    implicitness_theta = BoundedFloat(default_value=0.5, bounds=[0.5, 1.0], help='implicitness parameter theta. Value 0.5 implies Crank-Nicolson scheme, 1.0 implies fully implicit formulation.')


class ExplicitTimestepperOptions(TimeStepperOptions):
    use_automatic_timestep = Bool(True, help='Set time step automatically based on local CFL conditions.')

class GLSModelOptions(FrozenHasTraits):
    """Class for all time stepper options"""
    pass

class TimeStepperOptionsTrait(TraitType):
    default_value = ExplicitTimestepperOptions()
    info_text = 'a TimeStepperOptions instance'

    def validate(self, obj, value):
        if isinstance(value, TimeStepperOptions):
            return value
        self.error(obj, value)

class SubOptions(FrozenHasTraits):

    @observe('name')
    def _observe_name(self, change):
        old = change['old']
        newval = change['new']
        for target_trait, choices in self.children:
            if newval in choices:
                self.__setattr__(target_trait, choices[newval])

    def __init__(self, valid_values=[], default_value=None, help=None, children=None):
        name = Enum(valid_values,
                    default_value=default_value,
                    help=help).tag(config=True)
        self.add_traits(name=name)

        self.children = children
        # check completeness of children
        for target_trait, choices in self.children:
            for value in valid_values:
                assert value in choices, 'Cannot find key "{:}" in "{:}" dictionary'.format(value, target_trait)

        # set defaults
        assert default_value is not None
        for target_trait, choices in self.children:
            trait = choices[default_value]
            self.__setattr__(target_trait, trait)

        self._isfrozen = True


def attach_paired_options(options_name_trait, options_value_trait, default_values):
    """Attach paired options to a Configurable object.

    :arg options_name_trait: a tuple (name, Trait) for the options name (a choice).
    :arg options_value_trait: a tuple (name, Trait) for the options value (some configurable object)
    :arg default_values: a dict mapping valid values for the options name to valid defaults for the options value."""

    name, name_trait = options_name_trait
    value, value_trait = options_value_trait

    def _observer(self, change):
        "Observer called when the choice option is updated."
        setattr(self, value, default_values[change["new"]])

    def _default(self):
        "Dynamic default value setter"
        if hasattr(name_trait, 'default_value') and name_trait.default_value is not None:
            return default_values[name_trait.default_value]

    obs_handler = ObserveHandler(name, type="change")
    def_handler = DefaultHandler(value)

    def update_class(cls):
        "Programmatically update the class"
        # Set the new class attributes
        setattr(cls, name, name_trait)
        setattr(cls, value, value_trait)
        setattr(cls, "_%s_observer" % name, obs_handler(_observer))
        setattr(cls, "_%s_default" % value, def_handler(_default))
        # Mimic the magic metaclass voodoo.
        name_trait.class_init(cls, name)
        value_trait.class_init(cls, value)
        obs_handler.class_init(cls, "_%s_observer" % name)
        def_handler.class_init(cls, "_%s_default" % value)

        return cls
    return update_class


@attach_paired_options(("timestepper_type",
                        Enum(['SSPRK22', 'CrankNicolson'],
                             default_value='SSPRK22',
                             help='Name of the time integrator').tag(config=True)
                        ), ("timestepper_options",
                           Instance(TimeStepperOptions, args=()).tag(config=True)),
                        {
                            'CrankNicolson': CrankNicolsonOptions(),
                            'SSPRK22': ExplicitTimestepperOptions(),
                        })
@attach_paired_options(("turbulence_model_type",
                        Enum(['gls'],
                             default_value='gls',
                             help='Type of vertical turbulence model').tag(config=True)
                        ), ("gls_options",
                           Instance(GLSModelOptions, args=()).tag(config=True)),
                        {
                            'CrankNicolson': CrankNicolsonOptions(),
                            'SSPRK22': ExplicitTimestepperOptions(),
                        })
class ModelOptions(FrozenConfigurable):
    polynomial_degree = PositiveInteger(1, help='Polynomial degree of elements').tag(config=True)
    element_family = Enum(
        ['dg-dg', 'rt-dg', 'dg-cg'],
        default_value='dg-dg',
        help="""Finite element family

        2D solver supports 'dg-dg', 'rt-dg', or 'dg-cg' velocity-pressure pairs.
        3D solver supports 'dg-dg', or 'rt-dg' velocity-pressure pairs."""
        ).tag(config=True)

    use_nonlinear_equations = Bool(True, help='Use nonlinear shallow water equations').tag(config=True)
    solve_salinity = Bool(True, help='Solve salinity transport').tag(config=True)
    solve_temperature = Bool(True, help='Solve temperature transport').tag(config=True)
    use_implicit_vertical_diffusion = Bool(True, help='Solve vertical diffusion and viscosity implicitly').tag(config=True)
    use_bottom_friction = Bool(True, help='Apply log layer bottom stress in the 3D model').tag(config=True)
    use_parabolic_viscosity = Bool(
        False,
        help="""Use idealized parabolic eddy viscosity

        See :class:`.ParabolicViscosity`""").tag(config=True)
    use_grad_div_viscosity_term = Bool(
        False,
        help=r"""Include :math:`\nabla (\nu_h \nabla \cdot \bar{\textbf{u}})` term in the depth-averaged viscosity

        See :class:`.shallowwater_eq.HorizontalViscosityTerm` for details.""").tag(config=True)
    use_grad_depth_viscosity_term = Bool(
        True,
        help=r"""Include :math:`\nabla H` term in the depth-averaged viscosity

        See :class:`.shallowwater_eq.HorizontalViscosityTerm` for details.""").tag(config=True)

    use_ale_moving_mesh = Bool(
        True, help="Use ALE formulation where 3D mesh tracks free surface").tag(config=True)

    use_linearized_semi_implicit_2d = Bool(
        False, help="Use linearized semi-implicit time integration for the horizontal mode").tag(config=True)
    shallow_water_theta = BoundedFloat(
        0.5, bounds=[0.5, 1.0], help='Theta parameter for shallow water semi-implicit scheme').tag(config=True)
    use_turbulence = Bool(
        False, help="Activate turbulence model in the 3D model").tag(config=True)
    use_smooth_eddy_viscosity = Bool(
        False, help="Cast eddy viscosity to p1 space instead of p0").tag(config=True)

    use_turbulence_advection = Bool(
        False, help="Advect TKE and Psi in the GLS turbulence model").tag(config=True)
    use_baroclinic_formulation = Bool(
        False, help="Compute internal pressure gradient in momentum equation").tag(config=True)
    use_smagorinsky_viscosity = Bool(
        False, help="Use Smagorinsky horisontal viscosity parametrization").tag(config=True)
    smagorinsky_coefficient = FiredrakeConstant(
        Constant(0.1),
        help="""Smagorinsky viscosity coefficient :math:`C_S`

        See :class:`.SmagorinskyViscosity`.""").tag(config=True)

    use_limiter_for_tracers = Bool(
        False, help="Apply P1DG limiter for tracer fields").tag(config=True)
    use_lax_friedrichs_velocity = Bool(
        True, help="use Lax Friedrichs stabilisation in horizontal momentum advection.").tag(config=True)
    lax_friedrichs_velocity_scaling_factor = FiredrakeConstant(
        Constant(1.0), help="Scaling factor for Lax Friedrichs stabilisation term in horiozonal momentum advection.").tag(config=True)
    use_lax_friedrichs_tracer = Bool(
        True, help="Use Lax Friedrichs stabilisation in tracer advection.").tag(config=True)
    lax_friedrichs_tracer_scaling_factor = FiredrakeConstant(
        Constant(1.0), help="Scaling factor for tracer Lax Friedrichs stability term.").tag(config=True)
    check_volume_conservation_2d = Bool(
        False, help="""
        Compute volume of the 2D mode at every export

        2D volume is defined as the integral of the water elevation field.
        Prints deviation from the initial volume to stdout.
        """).tag(config=True)
    check_volume_conservation_3d = Bool(
        False, help="""
        Compute volume of the 3D domain at every export

        Prints deviation from the initial volume to stdout.
        """).tag(config=True)
    check_salinity_conservation = Bool(
        False, help="""
        Compute total salinity mass at every export

        Prints deviation from the initial mass to stdout.
        """).tag(config=True)
    check_salinity_overshoot = Bool(
        False, help="""
        Compute salinity overshoots at every export

        Prints overshoot values that exceed the initial range to stdout.
        """).tag(config=True)
    check_temperature_conservation = Bool(
        False, help="""
        Compute total temperature mass at every export

        Prints deviation from the initial mass to stdout.
        """).tag(config=True)
    check_temperature_overshoot = Bool(
        False, help="""
        Compute temperature overshoots at every export

        Prints overshoot values that exceed the initial range to stdout.
        """).tag(config=True)
    log_output = Bool(
        True, help="Redirect all output to log file in output directory").tag(config=True)
    timestep = PositiveFloat(
        10.0, help="Time step").tag(config=True)
    timestep_2d = PositiveFloat(
        10.0, help="""
        Time step of the 2d mode

        This option is only used in the 3d solver, if 2d mode is solved
        explicitly.
        """).tag(config=True)
    use_automatic_timestep = Bool(
        True, help="""
        Set time step automatically.

        Solver computes the largest stable time step based on user-defined
        velocity and viscosity scales. See
        :attr:`horizontal_velocity_scale`,
        :attr:`vertical_velocity_scale`,
        :attr:`horizontal_viscosity_scale`.
        """).tag(config=True)
    cfl_2d = PositiveFloat(
        1.0, help="Factor to scale the 2d time step OBSOLETE").tag(config=True)  # TODO OBSOLETE
    cfl_3d = PositiveFloat(
        1.0, help="Factor to scale the 2d time step OBSOLETE").tag(config=True)  # TODO OBSOLETE
    simulation_export_time = PositiveFloat(
        100.0, help="""
        Export interval in seconds

        All fields in fields_to_export list will be stored to disk and
        diagnostics will be computed
        """).tag(config=True)
    simulation_end_time = PositiveFloat(
        1000.0, help="Simulation duration in seconds").tag(config=True)
    horizontal_velocity_scale = FiredrakeConstant(
        Constant(0.1), help="""
        Maximum horizontal velocity magnitude

        Used to compute max stable advection time step.
        """).tag(config=True)
    vertical_velocity_scale = FiredrakeConstant(
        Constant(1e-4), help="""
        Maximum vertical velocity magnitude

        Used to compute max stable advection time step.
        """).tag(config=True)
    horizontal_viscosity_scale = FiredrakeConstant(
        Constant(1.0), help="""
        Maximum horizontal viscosity

        Used to compute max stable diffusion time step.
        """).tag(config=True)
    output_directory = Unicode(
        'outputs', help="Directory where model output files are stored").tag(config=True)
    no_exports = Bool(
        False, help="""
        Do not store any outputs to disk

        Disables VTK and HDF5 field outputs. and HDF5 diagnostic outputs.
        Used in CI test suite.
        """).tag(config=True)
    export_diagnostics = Bool(
        True, help="Store diagnostic variables to disk in HDF5 format").tag(config=True)
    use_quadratic_pressure = Bool(
        False, help="""
        Use P2DGxP2 space for baroclinic head.

        If element_family='dg-dg', P2DGxP1DG space is also used for the internal
        pressure gradient.

        This is useful to alleviate bathymetry-induced pressure gradient errors.
        If False, the baroclinic head is in the tracer space, and internal
        pressure gradient is in the velocity space.
        """).tag(config=True)
    use_quadratic_density = Bool(
        False, help="""
        Water density is projected to P2DGxP2 space.

        This reduces pressure gradient errors associated with nonlinear
        equation of state.
        If False, density is computed point-wise in the tracer space.
        """).tag(config=True)
    fields_to_export = List(
        trait=Unicode,
        default_value=['elev_2d', 'uv_2d', 'uv_3d', 'w_3d'],
        help="Fields to export in VTK format").tag(config=True)
    fields_to_export_hdf5 = List(
        trait=Unicode,
        default_value=[],
        help="Fields to export in HDF5 format").tag(config=True)
    verbose = Integer(0, help="Verbosity level").tag(config=True)
    linear_drag_coefficient = FiredrakeCoefficient(
        None, allow_none=True, help=r"""
        2D linear drag parameter :math:`L`

        Bottom stress is :math:`\tau_b/\rho_0 = -L \mathbf{u} H`
        """).tag(config=True)
    quadratic_drag_coefficient = FiredrakeCoefficient(
        None, allow_none=True, help=r"""
        Dimensionless 2D quadratic drag parameter :math:`C_D`

        Bottom stress is :math:`\tau_b/\rho_0 = -C_D |\mathbf{u}|\mathbf{u}`
        """).tag(config=True)
    manning_drag_coefficient = FiredrakeCoefficient(
        None, allow_none=True, help=r"""
        Manning-Strickler 2D quadratic drag parameter :math:`\mu`

        Bottom stress is :math:`\tau_b/\rho_0 = -g \mu^2 |\mathbf{u}|\mathbf{u}/H^{1/3}`
        """).tag(config=True)
    use_wetting_and_drying = Bool(
        False, help=r"""bool: Turn on wetting and drying

        Uses the wetting and drying scheme from Karna et al (2011).
        If ``True``, one should also set :attr:`wetting_and_drying_alpha` to control the bathymetry displacement.
        """).tag(config=True)
    wetting_and_drying_alpha = FiredrakeConstant(
        Constant(0.5), help=r"""
        Coefficient: Wetting and drying parameter :math:`alpha`.

        Used in bathymetry displacement function that ensures positive water depths. Unit is meters.
        """).tag(config=True)
    horizontal_diffusivity = FiredrakeCoefficient(
        None, allow_none=True, help="Horizontal diffusivity for tracers").tag(config=True)
    vertical_diffusivity = FiredrakeCoefficient(
        None, allow_none=True, help="Vertical diffusivity for tracers").tag(config=True)
    horizontal_viscosity = FiredrakeCoefficient(
        None, allow_none=True, help="Horizontal viscosity").tag(config=True)
    vertical_viscosity = FiredrakeCoefficient(
        None, allow_none=True, help="Vertical viscosity").tag(config=True)
    coriolis_frequency = FiredrakeCoefficient(
        None, allow_none=True, help="2D Coriolis parameter").tag(config=True)
    wind_stress = FiredrakeCoefficient(
        None, allow_none=True, help="Stress at free surface (2D vector function)").tag(config=True)
    atmospheric_pressure = FiredrakeCoefficient(
        None, allow_none=True, help="Atmospheric pressure at free surface, in pascals").tag(config=True)
    momentum_source_2d = FiredrakeCoefficient(
        None, allow_none=True, help="Source term for 2D momentum equation").tag(config=True)
    momentum_source_3d = FiredrakeCoefficient(
        None, allow_none=True, help="Source term for 3D momentum equation").tag(config=True)
    volume_source_2d = FiredrakeCoefficient(
        None, allow_none=True, help="Source term for 2D continuity equation").tag(config=True)
    salinity_source_3d = FiredrakeCoefficient(
        None, allow_none=True, help="Source term for salinity equation").tag(config=True)
    temperature_source_3d = FiredrakeCoefficient(
        None, allow_none=True, help="Source term for temperature equation").tag(config=True)
    constant_temperature = FiredrakeConstant(
        Constant(10.0), help="Constant temperature if temperature is not solved").tag(config=True)
    constant_salinity = FiredrakeConstant(
        Constant(0.0), help="Constant salinity if salinity is not solved").tag(config=True)

options = ModelOptions()