#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This file is part of pyFresnel.
# Copyright (c) 2012, Robert Steed
# Author: Robert Steed (rjsteed@talk21.com)
# License: GPL
# last modified 18.12.2012
import numpy as N
sin,arcsin,cos,arccos=N.sin,N.arcsin,N.cos,N.arccos
import pyFresnelInit #need this to find modules in parent directory.
import pyFresnel.transfer_matrix as TM
sqrt,exp=N.sqrt,N.exp

eps0=8.8541E-12 #Farads/metres -vacuum permittivity.
m_e=9.1094e-31 #Kg - mass of electron
q=1.6022e-19 #C - unit charge.
c=299792458  #m/s - speed of light
pi=N.pi

class material(object):
    """Like a dielectric, plasma, quantum well etc"""
    def __init__(self):
        """This base class has a circular definition for n and epsilon. One 
        function must be overridden in the derived class!"""        
        pass
        
    def epsilon(self):
        return self.n()**2

    def n(self):
        return sqrt(self.epsilon())
        
    def __len__(self):
        pass
        
    def __add__(self,other): # this might not work once we start using the Claussius-Claussis relation!
        """Add two derived instances of classes derived from material"""
        def new_epsilon(self2):
            return self.epsilon()+other.epsilon()
        newmat=material()
        newmat.epsilon=new_epsilon.__get__(newmat,material) # binds function to instance
        return newmat

class Lorentz_model(material):
    """Simple model of an absorbing oscillator / transition.
    Frequencies - whether we use real or natural frequency doesn't matter as long as we are consistant! 
    Remember that there is a difference of 2pi between the two: w=2*pi*f
    Note that normally the equations for the plasma frequency will give a natural frequency but that otherwise
    will be interested in real frequencies."""
    def __init__(self,w,w0,y,wp,f,eps_b):
        #
        #So if these are real frequencies, the plasma frequency must be real too (and the normal equations give a natural value).
        self.w=w
        self.w0=w0
        self.y=y
        self.wp=wp
        self.f=f
        self.eps_b=eps_b
    
    def epsilon(self):
        w,w0,y,wp,f,eps_b=self.w,self.w0,self.y,self.wp,self.f,self.eps_b
        eps=eps_b*(1+wp**2*f/(w0**2-w**2-2j*y*w))
        return eps
        
    @staticmethod   
    def wp(N,meff,eps_b):
        """N (m**-3) charge density
        meff (fraction of m_e) effective mass
        eps_b (unitless) background dielectric"""
        return sqrt(N*q**2/(meff*m_e*eps0*eps_b))
 
    

    
freq=N.arange(0,6e12,5e9) #Frequency range (Hz) (REAL)
#Real vs Natural Frequencies, whether we are using real or natural frequencies is not
#important for the dielectric constant which is unitless but is important for the
#calculation of the phase shift which requires a natural frequency.

eps_b=10.0

Ls=[Lorentz_model(w=freq,w0=w0,y=15e10,wp=1.6e12,f=1.0,eps_b=eps_b) for w0 in 1.5e12,4e12]
#w0 - frequency of transition (Hz) (real)
#y - broadening of transition (~Half Width Half Maximum) (Hz) (real)
#wp - Plasma frequency (affects the strength of the transition) (real)
#f - oscillator strength (also affects the strength of the transition - factor due to quantum mechanics of the transition/oscillator). Can leave at 1.0.
#eps_b - background dielectric constant.
    
#Setup Transfer Matrix
filterlist = [TM.Layer_eps(eps_b,None),
        TM.Layer_eps(Ls[0].epsilon(),5e-6), #dielectric constant, thickness (m)
        TM.Layer_eps(eps_b,1e-4),
        TM.LayerUniaxial_eps(eps_b,Ls[1].epsilon(),5e-6),
        TM.Layer_eps(eps_b,1e-4),
        #TM.Layer_Uniaxial(Ls[1].epsilon(),eps_b,5e-6),
        #TM.Layer(eps_b,1e-4),
        TM.Layer_eps(eps_b,None)]
#
f1 = TM.Filter(filterlist,
                w=2*pi*freq, #Frequency (natural) (Hz)
                pol='TM', #polarisation: either 'TM' or 'TE'
                theta=pi/4) # angle of incidence (radians)
    


if __name__=="__main__":
    
    print f1
    #print repr(f1)
    #print len(f1)
    #print f1.calculate_M()[0]
    #print f1._lambda((f1[0],f1[-1]))
    #print f1.calculate_r_t()[0]
    w,R,T=f1.calculate_R_T()
    
    import pylab as pl
    pl.figure(1) #,figsize=(7,8))
    THz=freq*1e-12
    ax1=pl.subplot(111)
    #
    ax1.plot(THz,R,label="reflection")
    ax1.plot(THz,T,label="Transmission")
    #
    ax1.legend()
    ax1.set_title("A Uniaxial Transfer Matrix")
    ax1.set_xlabel("Frequency (real) (THz)")
    ax1.set_xlim((1,6))
    #
    pl.show()
    
    
