# Repository: https://gitlab.com/quantify-os/quantify-core
# Licensed according to the LICENCE file on the main branch
"""Models and fit functions to be used with the lmfit fitting framework."""
from __future__ import annotations

from typing import Tuple

import lmfit
import numpy as np


def get_model_common_doc() -> str:
    """
    Returns a common docstring to be used with custom fitting
    :class:`~lmfit.model.Model` s.

    .. admonition:: Usage example for a custom fitting model
        :class: dropdown, tip

        See the usage example at the end of the :class:`~ResonatorModel` source-code:

        .. literalinclude:: ../quantify_core/analysis/fitting_models.py
            :pyobject: ResonatorModel
    """
    return (
        lmfit.models.COMMON_INIT_DOC.replace(
            ":class:`Model`", ":class:`~lmfit.model.Model`"
        )
        .replace("\n    ", "\n")
        .replace(" : str", " : :obj:`str`")
        .replace("['x']", ":code:`['x']`")
        .replace(", optional", "")
        .replace(" optional", "")
        .replace("{'raise', 'propagate', 'omit'}", "")
    )


def get_guess_common_doc() -> str:
    """
    Returns a common docstring to be used for the :meth:`~lmfit.model.Model.guess`
    method of custom fitting :class:`~lmfit.model.Model` s.

    .. admonition:: Usage example for a custom fitting model
        :class: dropdown, tip

        See the usage example at the end of the :class:`~ResonatorModel` source-code:

        .. literalinclude:: ../quantify_core/analysis/fitting_models.py
            :pyobject: ResonatorModel
    """
    return (
        lmfit.models.COMMON_GUESS_DOC.replace(
            " : Parameters", " : :class:`~lmfit.parameter.Parameters`"
        )
        .replace("\n    ", "\n")
        .replace(" optional", "")
        .replace(" : array_like", " : :class:`~numpy.ndarray`")
    )


def mk_seealso(
    function_name: str,
    role: str = "func",
    prefix: str = "\n\n",
    module_location: str = ".",
) -> str:
    """
    Returns a sphinx `seealso` pointing to a function.

    Intended to be used for building custom fitting model docstrings.

    .. admonition:: Usage example for a custom fitting model
        :class: dropdown, tip

        See the usage example at the end of the :class:`~ResonatorModel` source-code:

        .. literalinclude:: ../quantify_core/analysis/fitting_models.py
            :pyobject: ResonatorModel

    Parameters
    ----------
    function_name
        name of the function to point to
    role
        a sphinx role, e.g. :code:`"func"`
    prefix
        string preceding the `seealso`
    module_location
        can be used to indicate a function outside this module, e.g.,
        :code:`my_module.submodule` which contains the function.

    Returns
    -------
    :
        resulting string
    """
    return f"{prefix}.. seealso:: :{role}:`~{module_location}{function_name}`\n"


# pylint: disable=too-many-arguments
# pylint: disable=invalid-name
def hanger_func_complex_SI(
    f: float,
    fr: float,
    Ql: float,
    Qe: float,
    A: float,
    theta: float,
    phi_v: float,
    phi_0: float,
    alpha: float = 1,
) -> complex:
    r"""
    This is the complex function for a hanger (lambda/4 resonator).

    Parameters
    ----------
    f:
        frequency
    fr:
        resonance frequency
    A:
        background transmission amplitude
    Ql:
        loaded quality factor of the resonator
    Qe:
        magnitude of extrinsic quality factor :code:`Qe = |Q_extrinsic|`
    theta:
        phase of extrinsic quality factor (in rad)
    phi_v:
        phase to account for propagation delay to sample
    phi_0:
        phase to account for propagation delay from sample
    alpha:
        slope of signal around the resonance

    Returns
    -------
    :
        complex valued transmission


    See eq. S4 from Bruno et al. (2015)
    `ArXiv:1502.04082 <https://arxiv.org/abs/1502.04082>`_.

    .. math::

        S_{21} = A \left(1+\alpha \frac{f-f_r}{f_r} \right)
        \left(1- \frac{\frac{Q_l}{|Q_e|}e^{i\theta} }{1+2iQ_l \frac{f-f_r}{f_r}} \right)
        e^{i (\phi_v f + \phi_0)}

    The loaded and extrinsic quality factors are related to the internal and coupled Q
    according to:

    .. math::

        \frac{1}{Q_l} = \frac{1}{Q_c}+\frac{1}{Q_i}

    and

    .. math::

        \frac{1}{Q_c} = \mathrm{Re}\left(\frac{1}{|Q_e|e^{-i\theta}}\right)

    """
    slope_corr = 1 + alpha * (f - fr) / fr

    hanger_contribution = 1 - Ql / Qe * np.exp(1j * theta) / (
        1 + 2.0j * Ql * (f - fr) / fr
    )

    propagation_delay_corr = np.exp(1j * (phi_v * f + phi_0))
    S21 = A * slope_corr * hanger_contribution * propagation_delay_corr

    return S21


def cos_func(
    x: float,
    frequency: float,
    amplitude: float,
    offset: float,
    phase: float = 0,
) -> float:
    r"""
    An oscillating cosine function:

    :math:`y = \mathrm{amplitude} \times \cos(2 \pi \times \mathrm{frequency} \times x + \mathrm{phase}) +  \mathrm{offset}`

    Parameters
    ----------
    x:
        The independent variable (time, for example)
    frequency:
        A generalized frequency (in units of inverse x)
    amplitude:
        Amplitude of the oscillation
    offset:
        Output signal vertical offset
    phase:
        Phase offset / rad

    Returns
    -------
    :
        Output signal magnitude
    """  # pylint: disable=line-too-long

    return amplitude * np.cos(2 * np.pi * frequency * x + phase) + offset


def exp_decay_func(
    t: float,
    tau: float,
    amplitude: float,
    offset: float,
    n_factor: float,
) -> float:
    r"""
    This is a general exponential decay function:

    :math:`y = \mathrm{amplitude} \times \exp\left(-(t/\tau)^\mathrm{n\_factor}\right) + \mathrm{offset}`

    Parameters
    ----------
    t:
        time
    tau:
        decay time
    amplitude:
        amplitude of the exponential decay
    offset:
        asymptote of the exponential decay, the value at t=infinity
    n_factor:
        exponential decay factor

    Returns
    -------
    :
        Output of exponential function as a float
    """  # pylint: disable=line-too-long
    return amplitude * np.exp(-((t / tau) ** n_factor)) + offset


def exp_damp_osc_func(
    t: float,
    tau: float,
    n_factor: float,
    frequency: float,
    phase: float,
    amplitude: float,
    offset: float,
):
    r"""
    A sinusoidal oscillation with an exponentially decaying envelope function:

    :math:`y = \mathrm{amplitude} \times \exp\left(-(t/\tau)^\mathrm{n\_factor}\right)(\cos(2\pi\mathrm{frequency}\times t + \mathrm{phase}) + \mathrm{oscillation_offset}) + \mathrm{exponential_offset}`

    Parameters
    ----------
    t:
        time
    tau:
        decay time
    n_factor:
        exponential decay factor
    frequency:
        frequency of the oscillation
    phase:
        phase of the oscillation
    amplitude:
        initial amplitude of the oscillation
    oscillation_offset:
        vertical offset of cosine oscillation relative to exponential asymptote
    exponential_offset:
        offset of exponential asymptote

    Returns
    -------
    :
        Output of decaying cosine function as a float
    """  # pylint: disable=line-too-long

    oscillation = amplitude * (np.cos(2 * np.pi * frequency * t + phase))
    exp_decay = np.exp(-((t / tau) ** n_factor))
    osc_decay = oscillation * exp_decay + offset
    return osc_decay


# This class is used a literal include in the docs so the pylint options are here
# pylint: disable=empty-docstring
# pylint: disable=abstract-method
# pylint: disable=too-few-public-methods
class CosineModel(lmfit.model.Model):
    """
    Exemplary lmfit model with a guess for a cosine.

    .. note::

        The :mod:`lmfit.models` module provides several fitting models that might fit
        your needs out of the box.
    """

    def __init__(self, *args, **kwargs):
        # pass in the model's equation
        super().__init__(cos_func, *args, **kwargs)

        # configure constraints that are independent from the data to be fitted
        self.set_param_hint("frequency", min=0, vary=True)  # enforce positive frequency
        self.set_param_hint("amplitude", min=0, vary=True)  # enforce positive amplitude
        self.set_param_hint("offset", vary=True)
        self.set_param_hint(
            "phase", vary=True, min=-np.pi, max=np.pi
        )  # enforce phase range

    # pylint: disable=missing-function-docstring
    def guess(self, data, x, **kws) -> lmfit.parameter.Parameters:
        """
        guess parameters based on the data

        Parameters
        ----------
        data: np.ndarray
            Data to fit to
        x: np.ndarray
            Independet variable
        """

        self.set_param_hint("offset", value=np.average(data))
        self.set_param_hint("amplitude", value=(np.max(data) - np.min(data)) / 2)

        # Guess frequency and phase using Fourier Transform
        freq_guess, phase_guess = fft_freq_phase_guess(data, x)
        phase_wrap = (phase_guess + np.pi) % (2 * np.pi) - np.pi
        self.set_param_hint("frequency", value=freq_guess)
        self.set_param_hint("phase", value=phase_wrap)

        params = self.make_params()
        return lmfit.models.update_param_vals(params, self.prefix, **kws)

    # Same design patter is used in lmfit.models to inherit common docstrings.
    # We adjust these common docstrings to our docs build pipeline
    __init__.__doc__ = get_model_common_doc() + mk_seealso("cos_func")
    guess.__doc__ = get_guess_common_doc()


class ResonatorModel(lmfit.model.Model):
    """
    Resonator model

    Implementation and design patterns inspired by the
    `complex resonator model example <https://lmfit.github.io/lmfit-py/examples/example_complex_resonator_model.html>`_
    (`lmfit` documentation).

    """  # pylint: disable=line-too-long

    # pylint: disable=empty-docstring
    # pylint: disable=abstract-method
    def __init__(self, *args, **kwargs):
        # pass in the defining equation so the user doesn't have to later.
        super().__init__(hanger_func_complex_SI, *args, **kwargs)
        self.set_param_hint("Ql", min=0)  # Enforce Q is positive
        self.set_param_hint("Qe", min=0)  # Enforce Q is positive

        # Internal and coupled quality factor can be derived from fitted params
        self.set_param_hint("Qi", expr="1./(1./Ql-1./Qe*cos(theta))", vary=False)
        self.set_param_hint("Qc", expr="Qe/cos(theta)", vary=False)

    # pylint: disable=too-many-locals
    # pylint: disable=missing-function-docstring
    def guess(self, data, **kws) -> lmfit.parameter.Parameters:
        f = kws.get("f", None)
        if f is None:
            return None

        argmin_s21 = np.abs(data).argmin()
        fmin = f.min()
        fmax = f.max()
        # guess that the resonance is the lowest point
        fr_guess = f[argmin_s21]
        # assume the user isn't trying to fit just a small part of a resonance curve.
        Q_min = 0.1 * (fr_guess / (fmax - fmin))
        delta_f = np.diff(f)  # assume f is sorted
        min_delta_f = delta_f[delta_f > 0].min()
        Q_max = (
            fr_guess / min_delta_f
        )  # assume data actually samples the resonance reasonably
        Q_guess = np.sqrt(Q_min * Q_max)  # geometric mean, why not?
        (phi_0_guess, phi_v_guess) = resonator_phase_guess(
            data, f
        )  # Come up with a guess for phase velocity

        self.set_param_hint("fr", value=fr_guess, min=fmin, max=fmax)
        self.set_param_hint("Ql", value=Q_guess * 1.01, min=Q_min, max=Q_max)
        self.set_param_hint("Qe", value=Q_guess * 0.99, min=0)
        self.set_param_hint("A", value=np.mean(abs(data)), min=0)

        # The parameters below need a proper guess.
        self.set_param_hint("theta", value=0, min=-np.pi / 2, max=np.pi / 2)
        self.set_param_hint("phi_0", value=phi_0_guess)
        self.set_param_hint("phi_v", value=phi_v_guess)
        self.set_param_hint("alpha", value=0, min=-1, max=1)

        params = self.make_params()
        return lmfit.models.update_param_vals(params, self.prefix, **kws)

    # Same design patter is used in lmfit.models
    __init__.__doc__ = get_model_common_doc() + mk_seealso("hanger_func_complex_SI")
    guess.__doc__ = get_guess_common_doc()


class ExpDecayModel(lmfit.model.Model):
    """
    Model for an exponential decay, such as a qubit T1 measurement.
    """

    # pylint: disable=empty-docstring
    # pylint: disable=abstract-method
    # pylint: disable=too-few-public-methods

    def __init__(self, *args, **kwargs):
        # pass in the defining equation so the user doesn't have to later.

        super().__init__(exp_decay_func, *args, **kwargs)
        self.set_param_hint("tau", min=0)  # Enforce T1 is positive

        self.set_param_hint("amplitude", vary=True)
        self.set_param_hint("offset", vary=True)
        self.set_param_hint("n_factor", expr="1", vary=False)

    # pylint: disable=missing-function-docstring
    def guess(self, data, **kws) -> lmfit.parameter.Parameters:
        delay = kws.get("delay", None)

        if delay is None:
            return None

        # To guess the upper amplitude and offset,
        # use the first and last values of the data
        self.set_param_hint("offset", value=data[-1])
        self.set_param_hint("amplitude", value=data[0] - data[-1])

        # The guess for tau is somewhere in the middle of the time range
        tau = np.median(delay)

        self.set_param_hint("tau", value=tau, min=0)

        params = self.make_params()
        return lmfit.models.update_param_vals(params, self.prefix, **kws)

    # Same design patter is used in lmfit.models
    __init__.__doc__ = get_model_common_doc() + mk_seealso("exp_decay_func")
    guess.__doc__ = get_guess_common_doc()


class RabiModel(lmfit.model.Model):
    r"""
    Model for a Rabi oscillation as a function of the microwave drive amplitude.
    Phase of oscillation is fixed at :math:`\pi` in order to ensure that the oscillation
    is at a minimum when the drive amplitude is 0.
    """

    # pylint: disable=empty-docstring
    # pylint: disable=abstract-method
    # pylint: disable=too-few-public-methods

    def __init__(self, *args, **kwargs):
        # pass in the defining equation so the user doesn't have to later.
        super().__init__(cos_func, *args, **kwargs)

        # Enforce oscillation frequency is positive
        self.set_param_hint("frequency", min=0)

        # Fix the phase at pi so that the ouput is at a minimum when x=0
        self.set_param_hint("phase", expr="3.141592653589793", vary=False)

        # Pi-pulse amplitude can be derived from the oscillation frequency
        self.set_param_hint("amp180", expr="1/(2*frequency)", vary=False)

    # pylint: disable=missing-function-docstring
    def guess(self, data, **kws) -> lmfit.parameter.Parameters:
        drive_amp = kws.get("drive_amp", None)
        if drive_amp is None:
            return None

        amp_guess = abs(max(data) - min(data)) / 2  # amp is positive by convention
        offs_guess = np.mean(data)

        (freq_guess, _) = fft_freq_phase_guess(data, drive_amp)

        self.set_param_hint("frequency", value=freq_guess, min=0)
        self.set_param_hint("amplitude", value=amp_guess, min=0)
        self.set_param_hint("offset", value=offs_guess)

        params = self.make_params()
        return lmfit.models.update_param_vals(params, self.prefix, **kws)

    # Same design patter is used in lmfit.models
    __init__.__doc__ = get_model_common_doc() + mk_seealso("cos_func")
    guess.__doc__ = get_guess_common_doc()


class DecayOscillationModel(lmfit.model.Model):
    r"""
    Model for a decaying oscillation which decays to a point with 0 offset from
    the centre of the of the oscillation (as in a Ramsey experiment, for example).
    """

    # pylint: disable=empty-docstring
    # pylint: disable=abstract-method
    # pylint: disable=too-few-public-methods

    def __init__(self, *args, **kwargs):
        # pass in the defining equation so the user doesn't have to later.
        super().__init__(exp_damp_osc_func, *args, **kwargs)

        # Enforce oscillation frequency is positive
        self.set_param_hint("frequency", min=0)
        # Enforce amplitude is positive
        self.set_param_hint("amplitude", min=0)
        # Enforce decay time is positive
        self.set_param_hint("tau", min=0)

        # Fix the n_factor at 1
        self.set_param_hint("n_factor", expr="1", vary=False)

    # pylint: disable=missing-function-docstring
    def guess(self, data, **kws) -> lmfit.parameter.Parameters:
        t = kws.get("t", None)
        if t is None:
            raise ValueError(
                'Time variable "t" must be specified in order to guess parameters'
            )
            return None

        amp_guess = abs(max(data) - min(data)) / 2  # amp is positive by convention
        exp_offs_guess = np.mean(data)
        tau_guess = 2 / 3 * np.max(t)

        (freq_guess, phase_guess) = fft_freq_phase_guess(data, t)

        self.set_param_hint("frequency", value=freq_guess, min=0)
        self.set_param_hint("amplitude", value=amp_guess, min=0)
        self.set_param_hint("offset", value=exp_offs_guess)
        self.set_param_hint("phase", value=phase_guess)
        self.set_param_hint("tau", value=tau_guess, min=0)

        params = self.make_params()
        return lmfit.models.update_param_vals(params, self.prefix, **kws)

    # Same design patter is used in lmfit.models
    __init__.__doc__ = get_model_common_doc() + mk_seealso("exp_damp_osc_func")
    guess.__doc__ = get_guess_common_doc()


def resonator_phase_guess(s21: np.ndarray, freq: np.ndarray) -> Tuple[float, float]:
    """
    Guesses the phase velocity in resonator spectroscopy,
    based on the median of all the differences between consecutive phases.

    Parameters
    ----------
    s21:
        Resonator S21 data
    freq:
        Frequency of the spectroscopy pulse

    Returns
    -------
    phi_0:
        Guess for the phase offset
    phi_v:
        Guess for the phase velocity
    """
    phase = np.angle(s21)

    med_diff = np.median(np.diff(phase))
    freq_step = np.median(np.diff(freq))

    phi_v = med_diff / freq_step

    phi_0 = phase[0] - phi_v * freq[0]

    return phi_0, phi_v


def fft_freq_phase_guess(data: np.ndarray, t: np.ndarray) -> Tuple[float, float]:
    """
    Guess for a cosine fit using FFT, only works for evenly spaced points.

    Parameters
    ----------
    data:
        Input data to FFT
    t:
        Independent variable (e.g. time)

    Returns
    -------
    freq_guess:
        Guess for the frequency of the cosine function
    ph_guess:
        Guess for the phase of the cosine function
    """

    # Only first half of array is used, because the second half contains the
    # negative frequecy components, and we want a positive frequency.
    power = np.fft.fft(data)[: len(data) // 2]
    freq = np.fft.fftfreq(len(data), t[1] - t[0])[: len(power)]
    power[0] = 0  # Removes DC component from fourier transform

    # Use absolute value of complex valued spectrum
    abs_power = np.abs(power)
    freq_guess = abs(freq[abs_power == max(abs_power)][0])
    # the condition data == max(data) can have several solutions
    #               (for example when discretization is visible)
    # to prevent errors we pick the first solution
    ph_guess = 2 * np.pi - (2 * np.pi * t[data == max(data)] * freq_guess)[0]

    return freq_guess, ph_guess
