# -*- coding: utf-8 -*-
# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2014 and later, Alexander J G Pitchford
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################

# @author: Alexander Pitchford
# @email1: agp1@aber.ac.uk
# @email2: alex.pitchford@gmail.com
# @organization: Aberystwyth University
# @supervisor: Daniel Burgarth

"""
Pulse generator - Generate pulses for the timeslots
Each class defines a gen_pulse function that produces a float array of
size num_tslots. Each class produces a differ type of pulse.
See the class and gen_pulse function descriptions for details
"""

import numpy as np

import qutip.logging as logging
logger = logging.get_logger()

import qutip.control.dynamics as dynamics
import qutip.control.errors as errors

def create_pulse_gen(pulse_type='RND', dyn=None, pulse_params=None):
    """
    Create and return a pulse generator object matching the given type.
    The pulse generators each produce a different type of pulse,
    see the gen_pulse function description for details.
    These are the random pulse options:
        RND - Independent random value in each timeslot
        RNDFOURIER - Fourier series with random coefficients
        RNDWAVES - Summation of random waves
        RNDWALK1 - Random change in amplitude each timeslot
        RNDWALK2 - Random change in amp gradient each timeslot
    These are the other non-periodic options:
        LIN - Linear, i.e. contant gradient over the time
        ZERO - special case of the LIN pulse, where the gradient is 0
    These are the periodic options
        SINE - Sine wave
        SQUARE - Square wave
        SAW - Saw tooth wave
        TRIANGLE - Triangular wave
    If a Dynamics object is passed in then this is used in instantiate
    the PulseGen, meaning that some timeslot and amplitude properties
    are copied over.
    """

    if pulse_type == 'RND':
        return PulseGenRandom(dyn, params=pulse_params)
    if pulse_type == 'RNDFOURIER':
        return PulseGenRndFourier(dyn, params=pulse_params)
    if pulse_type == 'RNDWAVES':
        return PulseGenRndWaves(dyn, params=pulse_params)
    if pulse_type == 'RNDWALK1':
        return PulseGenRndWalk1(dyn, params=pulse_params)
    if pulse_type == 'RNDWALK2':
        return PulseGenRndWalk2(dyn, params=pulse_params)
    elif pulse_type == 'LIN':
        return PulseGenLinear(dyn, params=pulse_params)
    elif pulse_type == 'ZERO':
        return PulseGenZero(dyn, params=pulse_params)
    elif pulse_type == 'SINE':
        return PulseGenSine(dyn, params=pulse_params)
    elif pulse_type == 'SQUARE':
        return PulseGenSquare(dyn, params=pulse_params)
    elif pulse_type == 'SAW':
        return PulseGenSaw(dyn, params=pulse_params)
    elif pulse_type == 'TRIANGLE':
        return PulseGenTriangle(dyn, params=pulse_params)
    elif pulse_type == 'GAUSSIAN':  
        return PulseGenGaussian(dyn, params=pulse_params)
    elif pulse_type == 'CRAB_FOURIER':
        return PulseGenCrabFourier(dyn, params=pulse_params)
    elif pulse_type == 'GAUSSIAN_EDGE':  
        return PulseGenGaussianEdge(dyn, params=pulse_params)
    else:
        raise ValueError("No option for pulse_type '{}'".format(pulse_type))


class PulseGen:
    """
    Pulse generator
    Base class for all Pulse generators
    The object can optionally be instantiated with a Dynamics object,
    in which case the timeslots and amplitude scaling and offset
    are copied from that.
    Otherwise the class can be used independently by setting:
    tau (array of timeslot durations)
    or
    num_tslots and pulse_time for equally spaced timeslots

    Attributes
    ----------
    num_tslots : integer
        Number of timeslots, aka timeslices
        (copied from Dynamics if given)

    pulse_time : float
        total duration of the pulse
        (copied from Dynamics.evo_time if given)

    scaling : float
        linear scaling applied to the pulse
        (copied from Dynamics.initial_ctrl_scaling if given)

    offset : float
        linear offset applied to the pulse
        (copied from Dynamics.initial_ctrl_offset if given)

    tau : array[num_tslots] of float
        Duration of each timeslot
        (copied from Dynamics if given)

    lbound : float
        Lower boundary for the pulse amplitudes
        Note that the scaling and offset attributes can be used to fully
        bound the pulse for all generators except some of the random ones
        This bound (if set) may result in additional shifting / scaling
        Default is -Inf

    ubound : float
     
   Upper boundary for the pulse amplitudes
        Note that the scaling and offset attributes can be used to fully
        bound the pulse for all generators except some of the random ones
        This bound (if set) may result in additional shifting / scaling
        Default is Inf

    periodic : boolean
        True if the pulse
 generator produces periodic pulses

    random : boolean
        True if the pulse generator produces random pulses

    log_level : integer
        level of messaging output from the logger.
        Options are attributes of qutip.logging,
        in decreasing levels of messaging, are:
        DEBUG_INTENSE, DEBUG_VERBOSE, DEBUG, INFO, WARN, ERROR, CRITICAL
        Anything WARN or above is effectively 'quiet' execution,
        assuming everything runs as expected.
        The default NOTSET implies that the level will be taken from
        the QuTiP settings file, which by default is WARN
    """
    def __init__(self, dyn=None, params=None):
        self.parent = dyn
        self.params = params
        self.reset()

    def reset(self):
        """
        reset attributes to default values
        """
        if isinstance(self.parent, dynamics.Dynamics):
            dyn = self.parent
            self.num_tslots = dyn.num_tslots
            self.pulse_time = dyn.evo_time
            self.scaling = dyn.initial_ctrl_scaling
            self.offset = dyn.initial_ctrl_offset
            self.tau = dyn.tau
            self.set_log_level(dyn.log_level)
        else:
            self.num_tslots = 100
            self.pulse_time = 1.0
            self.scaling = 1.0
            self.tau = None
            self.offset = 0.0
            self.log_level = logger.level

        self._uses_time = False
        self.time = None
        self._pulse_initialised = False
        self.periodic = False
        self.random = False
        self.lbound = None
        self.ubound = None
        self.ramping_pulse = N