# -*- coding: utf-8 -*-
"""
Created on Fri Feb 18 17:54:20 2022

@author: voumardt
"""
import numba
import numpy as np
from scipy.constants import c
import time as timing


class Solver():
    def __init__(self, model, z_pts=500, local_error=0.00001, adaptive_factor=1.1,
                 max_dz=None, method=None):
        """
        Solver class. Provides a stepper which performs integration of the PDE
        as well as error control and results saving.

        Parameters
        ----------
        model : Model object
            Physical model of the problem to be integrated.
        z_pts : int, optional
            Number of results to keep along the propagation axis. The default is 500.
        local_error : float, optional
            Tolerated local error. The default is 0.00001.
        adaptive_factor : float, optional
            Adaptive factor for the step size control. Bigger values lead to
            more aggressive step size control. The default is 1.1.
        max_dz : float, optional
            Maximum step size allowed. Particularly useful when simulating
            quasi-phase matched structures. The default is None.
        method: string
            Method used for integration. Can be RK4IP, ERK4IP or DP5IP.
            Default is DP5IP.
        """
        self.model = model
        self.z_pts = z_pts
        self.local_error = local_error
        self.adaptive_factor = adaptive_factor
        self.max_dz = max_dz
        self.z_save = np.linspace(0, self.model.waveguide.length, self.z_pts)
        self.method_name = self._set_method_name(method)
    
    def _set_method_name(self, method):
        """
        Helper function to choose an appropriate solver. If a method is not
        provided by the user, it  will look into the model used to see if
        there is a recommended solver. Otherwise, defaults to the model that
        should be mostly used, i.e. DP5IP.

        Parameters
        ----------
        method : string
            Name of the solver.

        Returns
        -------
        string
            Name of the solver.
        """
        if method is None:
            if hasattr(self.model, '_recommended_solver'):
                return self.model._recommended_solver
            else:
                return 'DP5IP'
        else:
            return method
        
    def _compute_error(self, field_1, field_2):
        """
        Error estimation for adaptive step size

        Parameters
        ----------
        field_1 : array
            Reference field.
        field_2 : array
            Test field.

        Returns
        -------
        float
            Error estimate between the two fields.
        """
        return numbanorm(field_1 - field_2)/numbanorm(field_1)
    
    def _dormand_prince_step(self, field, dz):
        """
        Dormand Prince order 5(4) adaptive scheme in the interaction picture. This
        new stepper uses the interaction picture in an embedded scheme of order 5.
        Tests have shown it to be up to four times faster than RK4IP.

        Parameters
        ----------
        field : array
            Current field.
        dz : float
            Current step size.

        Returns
        -------
        array
            New field estimate.
        
        float
            Error estimate.
        """
        if self._k_7 is None:
            k_1 = dz*self.model.dispersion_step(self.model.nonlinear_term(field), dz)
        else:
            k_1 = self.model.dispersion_step(self._k_7, dz)
        a_int = self.model.dispersion_step(field, dz)
        k_2 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + k_1/5, -4*dz/5)), 4*dz/5)
        k_3 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + 3*k_1/40 + 9*k_2/40, -7*dz/10)), 7*dz/10)
        k_4 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + 44*k_1/45 - 56*k_2/15 + 32*k_3/9, -dz/5)), dz/5)
        k_5 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + 19372*k_1/6561 - 25360*k_2/2187 + 64448*k_3/6561 - 212*k_4/729, -dz/9)), dz/9)
        k_6 = dz*self.model.nonlinear_term(a_int + 9017*k_1/3168 - 355*k_2/33 + 46732*k_3/5247 + 49*k_4/176 - 5103*k_5/18656)
        sol_5 = a_int + 35*k_1/384 + 500*k_3/1113 + 125*k_4/192 - 2187*k_5/6784 + 11*k_6/84
        self._k_7 = dz*self.model.nonlinear_term(sol_5)
        sol_4 = a_int + 5179*k_1/57600 + 7571*k_3/16695 + 393*k_4/640 - 92097*k_5/339200 + 187*k_6/2100 + self._k_7/40
        error = self._compute_error(sol_5, sol_4)
        return sol_5, error
    
    def _rk546m(self, field, dz):
        """
        Some other Dormand-Prince method that has one less step but does not
        leverage the interaction picture and does not have the FSAL property.
        Usually worse than the other dormand-prince method.
        References:
            DOI: 10.1016/0771-050X(80)90013-3

        Parameters
        ----------
        field : array
            Current field.
        dz : float
            Current step size.

        Returns
        -------
        array
            New field estimate.
        
        float
            Error estimate.
        """
        k_1 = dz*self.model.nonlinear_term(field)
        k_2 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + k_1/5, dz/5)), -dz/5)
        k_3 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + 3*k_1/40 + 9*k_2/40, 3*dz/10)), -3*dz/10)
        k_4 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + 3*k_1/10 - 9*k_2/10 + 6*k_3/5, 3*dz/5)), -3*dz/5)
        k_5 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + 226*k_1/729 - 25*k_2/27 + 880*k_3/729 + 55*k_4/729, 2*dz/3)), -2*dz/3)
        k_6 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field  - 181*k_1/270 + 5*k_2/2 - 266*k_3/297 - 91*k_4/27 + 189*k_5/55, dz)), -dz)
        sol_5 = self.model.dispersion_step(field + 19*k_1/216 + 1000*k_3/2079 - 125*k_4/216 + 81*k_5/88 + 5*k_6/56, dz)
        sol_4 = self.model.dispersion_step(field + 31*k_1/540 + 190*k_3/297 - 145*k_4/108 + 351*k_5/220 + k_6/20, dz)
        error = self._compute_error(sol_5, sol_4)
        return sol_5, error
    
    def _adams_bashforth54_step(self, field, dz):
        """To be implemented"""
        pass

    def _rk4ip_single_step(self, field, dz):
        """
        Runge-Kutta order 4 step in the interaction picture (RK4IP)

        Parameters
        ----------
        field : array
            Current field.
        dz : float
            Current step size.

        Returns
        -------
        array
            New field estimate.
        """
        dispersion_op = numbaexp(-1j*self.model.waveguide.k*dz/2)
        a_int = dispersion_op*field
        k_1 = dz*dispersion_op*self.model.nonlinear_term(field)
        k_2 = dz*self.model.nonlinear_term(a_int + k_1/2)
        k_3 = dz*self.model.nonlinear_term(a_int + k_2/2)
        k_4 = dz*self.model.nonlinear_term(dispersion_op*(a_int + k_3))
        return dispersion_op*(a_int + k_1/6 + k_2/3 + k_3/3) + k_4/6
    
    def _rk4ip_step(self, field, dz):
        """
        Runge-Kutta order 4 adaptive scheme in the interaction picture (RK4IP)
        with step doubling and combining. Provided for comparison but slower
        than the dormand-prince method.
        References:
            DOI: 10.1109/JLT.2007.909373
            DOI: 10.1109/JLT.2009.2021538

        Parameters
        ----------
        field : array
            Current field.
        dz : float
            Current step size.

        Returns
        -------
        array
            New field estimate.
        
        float
            Error estimate.
        """
        sol_single = self._rk4ip_single_step(field, dz)
        sol_double = self._rk4ip_single_step(field, dz/2)
        sol_double = self._rk4ip_single_step(sol_double, dz/2)
        error = self._compute_error(sol_double, sol_single)
        sol_5 = (sol_double*16 - sol_single)/15
        return sol_5, error
    
    def _erk4ip_step(self, field, dz):
        """
        Embedded Runge-Kutta order 4(3) adaptive scheme in the interaction
        picture (ERK4IP). Can be faster than DP5IP if accuracy is irrelevant,
        for example in SPM only cases. Otherwise, DP5IP performs better.
        References:
            DOI: 10.1016/j.cpc.2012.12.020

        Parameters
        ----------
        field : array
            Current field.
        dz : float
            Current step size.

        Returns
        -------
        array
            New field estimate.
        
        float
            Error estimate.
        """
        dispersion_op = numbaexp(-1j*self.model.waveguide.k*dz/2)
        a_int = dispersion_op*field
        if self._k_5 is None:
            k_1 = dz*dispersion_op*self.model.nonlinear_term(field)
        else:
            k_1 = dispersion_op*self._k_5
        k_2 = dz*self.model.nonlinear_term(a_int + k_1/2)
        k_3 = dz*self.model.nonlinear_term(a_int + k_2/2)
        k_4 = dz*self.model.nonlinear_term(dispersion_op*(a_int + k_3))
        beta = dispersion_op*(a_int + k_1/6 + k_2/3 + k_3/3)
        sol_4 = beta + k_4/6
        self._k_5 = dz*self.model.nonlinear_term(sol_4)
        sol_3 = beta + k_4/15 + self._k_5/10
        error = self._compute_error(sol_4, sol_3)
        return sol_4, error
    
    def _integrate_to_z(self, z):
        """
        Integrate to z with adaptive step size.

        Parameters
        -------
        z : float
            Target propagation length.
        """
        while self.stock_z < z:
            # Actuate position inside the waveguide
            self.model.waveguide.z = self.stock_z
            
            # Step of dz
            vec_eval, error = self.step(self.stock_vec, self.dz)
            
            # Error control
            if error >= 2*self.local_error:
                self.dz /= 2
                print('Discarded')
                if self.method_name == 'DP5IP':
                    self._k_7 = None
                if self.method_name == 'ERK4IP':
                    self._k_5 = None
                continue
            elif self.local_error <= error < 2*self.local_error:
                self.stock_z += self.dz
                self.dz /= self.adaptive_factor
            elif 0.5*self.local_error <= error < self.local_error:
                self.stock_z += self.dz
            else:
                self.stock_z += self.dz
                if self.max_dz is None or self.dz < self.max_dz:
                    self.dz *= self.adaptive_factor
            self.stock_vec = vec_eval
    
    def solve(self):
        """
        Integrate over waveguide length with adaptive step size
        """
        # Initialization
        self.spectrum = np.zeros((self.z_pts, self.model.waveguide.t_pts), dtype='complex128')
        self.dz = np.mean(np.diff(self.z_save))*np.sqrt(np.sqrt(self.local_error))
        self.stock_z = 0
        self.stock_vec = np.fft.fft(self.model.light.field_t_in)
        if self.method_name == 'DP5IP':
            self._k_7 = None
            self.step = self._dormand_prince_step
        elif self.method_name == 'ERK4IP':
            self._k_5 = None
            self.step = self._erk4ip_step
        elif self.method_name == 'RK4IP':
            self.step = self._rk4ip_step
        elif self.method_name == 'RK546M':
            self.step = self._rk546m
        timer = timing.perf_counter()
            
        # Total integration
        for i, z in enumerate(self.z_save):
            print('z: {}/{}'.format(i+1, self.z_pts))
            
            # Integrate over dz
            self._integrate_to_z(z)
            
            # Save spectrum
            self.spectrum[i, :] = np.fft.fftshift(self.stock_vec)\
                *numbaexp(1j*(self.model.waveguide.betas[0] + self.model.waveguide.betas[1]*self.model.waveguide.rel_omega)*self.stock_z)
        
        # Save waveform
        print(f"{timing.perf_counter() - timer} seconds for integration.")                
        self.waveform = np.fft.ifft(np.fft.ifftshift(self.spectrum, axes=1), axis=1)
        self.model.light._set_propagation(self.z_save, self.model.waveguide.freq, self.spectrum, self.waveform)
        

"""
Optimized functions
"""
### Numba exponential
@numba.njit([numba.complex128[:](numba.complex128[:]),numba.complex64[:](numba.complex64[:])])
def numbaexp(x):
    return np.exp(x)

### Numba norm
@numba.njit([numba.float64(numba.complex128[:])])
def numbanorm(field):
    return np.linalg.norm(field)
