# -*- coding: utf-8 -*-
"""
Created on Tue Apr 25 19:12:18 2017

@author: pablo
"""

import os 

dir_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

import sys
 
sys.path.append(dir_path) #change to src
from src import np,math,interp1d,Hyperplume,griddata,interp2d,plt

import matplotlib.lines as mlines



class AEM(Hyperplume):
    
    """Asymptotic Expansion Model of a plasma plume expansion.Class AEM inherits methods  from 
    parent class Hyperplume, and particularizes them.All initial inputs must be given in dimensional form.
    
    """
    
    def __init__(self,plasma={'Electrons': {'Gamma': 1,'T_0_electron': 2.1801714e-19,'q_electron': -1.6e-19},'Ions': {'mass_ion': 2.1801714e-25, 'q_ion': 1.6e-19}},z_span=np.linspace(0,100,500),r_span=np.linspace(0,40,500),n_init=np.linspace(1,100,500),uz_init=np.linspace(1,100,500),ur_init=np.linspace(0,100,500),sol_order=0):
        
        """ Class method __init__ is used as class constructor. Calls parent class Hyperplume constructor method __init__ to 
        store main plasma properties as attributes in the class.
        
        Args:
            plasma (dict): simple_plasma object dictionary containing basic plasma parameters.
            z_span (numpy.ndarray): axial region where the problem will be integrated.
            r_span (numpy.ndarray): initial far-field plasma radial profile.
            n_init (numpy.ndarray): initial far-field plum density front.
            uz_init (numpy.ndarray): initial far-region plume axial velocity profile
            ur_init (numpy.ndarray): initial fr-region plume radial velocity profile
            sol_order (int): Integer defining the AEM correction order for the plume integration.
                            -0: AEM "cold beam" zeroth order solution
                            -1: AEM first order correction
                            -2: Second Order Correction
                            
        Usage:
            >>>>>> e_charge,ion_mass,Plasma_temp,gamma_value=1.6e-19,2.1801714e-25,2.1801714e-19,1 #Main Plasma Parameters
            >>> Plasma = Hyperplume().simple_plasma(e_charge,ion_mass,Plasma_temp,gamma_value) #Loading Plasma dict
            >>> z_span = np.linspace(0,110,5000) # Axial plume grid for integration
            >>> r_0 = np.linspace(0,10,5000) #Initial plume radial profile
            >>> n0 =  np.exp(-6.15/2 * r_0**2) #Initial far-field plume density
            >>> uz0,,ur0  = np.linspace(20000,20000,100),np.linspace(0,40000,100) #Intial far-field plume axial and radial velocity fronts
            >>> AEM_Order = 2 # AEM model solution
            >>> PlumeAEM = AEM(Plasma,z_span,eta_0,n0,uz0,ur0,AEM_Order) #Creation of AEM plume
        
        Other important class attributes loaded in the AEM constructor are:
            d0 (numpy.ndarray): far field initial divergence ur0/uz0.
            d0p (numpy.ndarray): derivative of plume initial divergence
            eps (float): AEM expansion parameter 1/M_{0}^{2}
            uz0p (numpy.ndarray): derivative of initial far region axial velocity
            duz0p (numpy.ndarray): derivative of initial far region radial velocity
            z_grid,r_grid (numpy.ndarray): Plume grids where AEM problem is integrated
            
        To access these attributes, for instance:
            >>> print(PlumeAEM.d0)
            >>> print(PlumeAEM.eps)
        """
        
        #Call parent class Hyperplume constructor method to store main plasma properties as attributes in the AEM class."""
        super(AEM,self).__init__(plasma,z_span,r_span,n_init)
        
        self.uz0,self.ur0,self.d0 = uz_init,ur_init,ur_init/uz_init #Load additional AEM plasma plume properties
        self.alpha0=math.degrees(math.atan(interp1d(self.eta,self.d0)(1))) #Initial Plume divergence at the 95% streamline
        self.order = sol_order #AEM Solution Order
        self.eps = self.Gamma*self.T_0/(self.m_ion*(uz_init[0]**2 + ur_init[0]**2)) #residual expansion parameter. Inverse of squared Mach Number 
        self.M0 = np.sqrt(1/self.eps) #Plume Mach Nmeber
        
        "Derivatives of initial front"
        
        self.d0p = self.eta_deriver(self.eta,self.d0) #derivative of plume divergence
        self.d0p[0],self.d0p[-1] = self.d0[1]/self.eta[1],self.d0p[-2] + (self.d0p[-2] - self.d0p[-3]) #Edge vetor prime conditions
        self.uz0p = self.eta_deriver(self.eta,self.uz0) #derivative of initial axial velocity
        self.uz0p[0],self.uz0p[-1] = 0,self.uz0p[-2] + (self.uz0p[-2] - self.uz0p[-3])
        self.duz0p = self.eta_deriver(self.eta,self.d0*self.uz0) #derivatie of initial radial velocity
        self.duz0p[0],self.duz0p[-1] = self.duz0p[1]/self.eta[1],self.duz0p[-2] + (self.duz0p[-2] - self.duz0p[-3])
        
        "Grid Set Up"
    
        self.z_grid,self.r_grid = self.grid_setup(self.z_span.size,self.eta.size) #2D grids of z, and r points in the plume
                
    def solver(self):
        
        """ Class method Solver integrates the AEM model equations in the specified plume grid. The method stores the different order
        plasma properties in matrixes of size (mxn), where m,n are the number of z,r points,respectively.Porperties such as
        density,temperature,electric field,etc are calculated and saved as attributes of the class in this matrix form.
        
        Usage:
            >>> PlumeAEM = AEM(Plasma,z_span,eta_0,n0,uz0,ur0,AEM_Order) #Creation of AEM plume
            >>> PlumeAEM.solver() # be sure to create a valid AEM plasma plume before applying the plume solver method
        
        Main Plume properties solved and saved by the method as class attributes:
            lnn (numpy.ndarray): 3-D matrix containing density values (logarithmic) for the three possible AEM solution orders.
            uz (numpy.ndarray): 3-D matrix containing axial velocity values for the three possible AEM solution orders.
            ur (numpy.ndarray): 3-D matrix containing radial velocity values for the three possible AEM solution orders.
            T (numpy.ndarray): 3-D matrix containing plasma Temperature values for the three possible AEM solution orders.
            phi (numpy.ndarray):3-D matrix containing plasma ambipolar electric field for the three possible AEM solution orders.
            div (numpy.ndarray): 3-D matrix containing plume divergence values for the three possible AEM solution orders.
            eta_ (int,numpy.ndarray): 3-D matrix containing ion current streamlines for the three possible AEM solution orders.
        
        To access these varibles,for instance:
            >>> PlumeAEM.lnn[0,:,:] #density values for Cold Beam Zeroth Order AEM solution of plume expansion in the grid
            >>> PlumeAEM.uz[1,:,:] # axial velocity First Order AEM solution
            >>> PlumeAEM.T[2,:,:] ## Temperature values for Second Order AEM solution
        """
        
        self.__zpts = np.shape(self.z_grid)[1] #Number of axial steps