# -*- coding: utf-8 -*-
"""
Created on Thu Jun 04 13:53:39 2015
This file is part of pyNLO.

    pyNLO is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    pyNLO is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with pyNLO.  If not, see <http://www.gnu.org/licenses/>.
@author: ycasg
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

import matplotlib.pyplot as plt
from numpy.fft import fftshift, ifftshift
from pynlo.interactions.FourWaveMixing import global_variables
from pynlo.light.PulseBase import Pulse
import gc
import scipy.fftpack


try:
    import pyfftw
    PYFFTW_AVAILABLE=True
except:
    PYFFTW_AVAILABLE=False
    

class SSFM:
    METHOD_SSFM,METHOD_RK4IP = range(2)    
    def __init__(self,  local_error = 0.001, dz = 1e-5,
                 disable_Raman = False, disable_self_steepening = False,
                 suppress_iteration = True, USE_SIMPLE_RAMAN = False,
                 f_R = 0.18, f_R0 = 0.18, tau_1 = 0.0122, tau_2 = 0.0320):
        """
        This initialization function sets up the parameters of the SSFM.
        """
        self.iter = 0
        self.last_h = -1.0
        self.last_dir = 0.0
        self.eta = 5
        self.local_error = local_error
        self.method = SSFM.METHOD_RK4IP
        self.disable_Raman = disable_Raman
        self.disable_self_steepening = disable_self_steepening
        self.USE_SIMPLE_RAMAN = USE_SIMPLE_RAMAN

        # Raman fraction; may change depending upon which calculation method 
        # is used 
        self.f_R = f_R
        # The value for the old-style Raman response
        self.f_R0 = f_R0
        
        self.tau_1 = tau_1
        self.tau_2 = tau_2
        self.dz = dz
        self.dz_min = 1e-12
        self.suppress_iteration = suppress_iteration


    def load_fiber_parameters(self, pulse_in, fiber, output_power, z=0):
        """
        This funciton loads the fiber parameters into class variables.
        """
        self.betas[:]  =  fiber.get_betas(pulse_in, z=z)
        # self.alpha[:]  = -fiber.get_gain(pulse_in, output_power) # currently alpha cannot change with z
        self.gamma     =  fiber.get_gamma(z=z)
        
        self.betas[:]  = self.conditional_fftshift(self.betas)
        # self.alpha[:]   = self.conditional_fftshift(self.alpha)
        


    def setup_fftw(self, pulse_in, fiber, output_power, raman_plots = False):
        ''' Call immediately before starting Propagate. This function does two
        things:\n
        1) it sets up byte aligned arrays for fftw\n
        2) it fftshifts betas, omegas, and the Raman response so that no further\n
            shifts are required during integration. This saves lots of time.'''

        
        self.n = pulse_in.NPTS
        
        if PYFFTW_AVAILABLE:
            
            self.fft_input    = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.fft_output   = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.ifft_input   = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.ifft_output  = pyfftw.empty_aligned(self.n, dtype='complex128')      
    
            self.fft_input_2  = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.fft_output_2 = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.ifft_input_2 = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.ifft_output_2= pyfftw.empty_aligned(self.n, dtype='complex128')      

            self.ifft_input_3 = pyfftw.empty_aligned(self.n, dtype='complex128')      
            self.ifft_output_3= pyfftw.empty_aligned(self.n, dtype='complex128')      
    
    
            # To be double sure that there are no problems, also make 2 copies of
            # the FFT objects. This lets us nest ifft_2 around a function using ifft
            # without worrying about potential problems.
            #self.fft    = pyfftw.builders.fft(self.fft_input)
            #self.fft_2  = pyfftw.builders.fft(self.fft_input_2)
            #self.ifft   = pyfftw.builders.fft(self.ifft_input)
            #self.ifft_2 = pyfftw.builders.fft(self.ifft_input_2)
            self.fft = pyfftw.FFTW(self.fft_input,
                                   self.fft_output,
                                   direction='FFTW_BACKWARD')
            self.fft_2 = pyfftw.FFTW(se
lf.fft_input_2,
                                     self.fft_output_2,
                                     direction='FFTW_BACKWARD')
            
            self.ifft = pyfftw.FFTW(self.ifft_input,
                                    self.ifft_output,
                                    direction='FFTW_FORWARD')
           
 self.ifft_2 = pyfftw.FFTW(self.ifft_input_2,
                                      self.ifft_output_2,
                                      direction='FFTW_FORWARD')
            self.ifft_3 = pyfftw.FFTW(self.ifft_input_3,
                                      self.ifft_output_3,
                                      direction='FFTW_FORWARD')                                      
            
        else:
            self.fft_input    = np.ndarray((self.n,), dtype='complex128')
            self.fft_output   = np.ndarray((self.n,), dtype='complex128')
            self.ifft_input   = np.ndarray((self.n,), dtype='complex128')
            self.ifft_output  = np.ndarray((self.n,), dtype='complex128')

            self.fft_input_2  = np.ndarray((self.n,), dtype='complex128')
            self.fft_output_2 = np.ndarray((self.n,), dtype='complex128')
            self.ifft_input_2 = np.ndarray((self.n,), dtype='complex128')
            self.ifft_output_2= np.ndarray((self.n,), dtype='complex128')
                                 
            # self.fft_input    = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            # self.fft_output   = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            # self.ifft_input   = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            # self.ifft_output  = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            #
            # self.fft_input_2  = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            # self.fft_output_2 = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            # self.ifft_input_2 = pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
            # self.ifft_output_2= pyfftw.empty_aligned(self.n, fft_n, dtype='complex128')
    
            
        self.A_I    = np.ndarray((self.n,), dtype='complex128')      
        
        self.A2     = np.ndarray((self.n,), dtype='complex128')
        self.exp_D  = np.ndarray((self.n,), dtype='complex128')      
        self.k1     = np.ndarray((self.n,), dtype='complex128')      
        self.k2     = np.ndarray((self.n,), dtype='complex128')      
        self.k3     = np.ndarray((self.n,), dtype='complex128')      
        self.k4     = np.ndarray((self.n,), dtype='complex128')      
        self.temp   = np.ndarray((self.n,), dtype='complex128')      
        self.Aw     = np.ndarray((self.n,), dtype='complex128')      
        self.A2w    = np.ndarray((self.n,), dtype='complex128')      
        self.dA     = np.ndarray((self.n,), dtype='complex128')      
        self.dA2    = np.ndarray((self.n,), dtype='complex128')      
        self.R_A2   = np.ndarray((self.n,), dtype='complex128')      
        self.dR_A2  = np.ndarray((self.n,), dtype='complex128')      
        self.omegas = np.ndarray((self.n,), dtype='complex128')              
        self.alpha  = np.ndarray((self.n,), dtype='complex128')      
        self.betas  = np.ndarray((self.n,), dtype