#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 18 10:02:39 2023

@author: victorcolas
"""

import numpy as np
import matplotlib.pyplot as plt


def CoeffFresnelTE(ni,nt,thetai):
    if thetai>=np.arcsin(nt/ni):
        thetat=np.pi/2
    else:
        thetat=np.arcsin(ni/nt*np.sin(thetai))
    r_TE=(ni*np.cos(thetai)-nt*np.cos(thetat))/(ni*np.cos(thetai)+nt*np.cos(thetat))
    t_TE=(2*ni*np.cos(thetai))/(ni*np.cos(thetai)+nt*np.cos(thetat))
    return r_TE,t_TE

def CoeffFresnelTM(ni,nt,thetai):
    if thetai>=np.arcsin(nt/ni):
        thetat=np.pi/2
    else:
        thetat=np.arcsin(ni/nt*np.sin(thetai))
    r_TE=(nt*np.cos(thetai)-ni*np.cos(thetat))/(nt*np.cos(thetai)+ni*np.cos(thetat))
    t_TE=(2*ni*np.cos(thetai))/(nt*np.cos(thetai)+ni*np.cos(thetat))
    return r_TE,t_TE


ni=1.
nt=1.5
n=100

listeThetaInc=np.linspace(0,np.pi/2,n)
liste_R_TE=np.zeros(n)
liste_r_TE=np.zeros(n)
liste_T_TE=np.zeros(n)
liste_t_TE=np.zeros(n)
liste_R_TM=np.zeros(n)
liste_r_TM=np.zeros(n)
liste_T_TM=np.zeros(n)
liste_t_TM=np.zeros(n)

for k in range(0,n):
    thetai=listeThetaInc[k]
    # TE
    r_TE,t_TE=CoeffFresnelTE(ni,nt,thetai)
    liste_r_TE[k]=r_TE
    liste_R_TE[k]=r_TE**2
    liste_t_TE[k]=t_TE
    liste_T_TE[k]=1-r_TE**2
    # TM
    r_TM,t_TM=CoeffFresnelTM(ni,nt,thetai)
    liste_r_TM[k]=r_TM
    liste_R_TM[k]=r_TM**2
    liste_t_TM[k]=t_TM
    liste_T_TM[k]=1-r_TM**2
 
# En amplitude 
plt.figure()
plt.plot(listeThetaInc*180/np.pi,liste_r_TE,label="r_TE",color="blue",ls='solid')
plt.plot(listeThetaInc*180/np.pi,liste_t_TE,label="t_TE",color="blue",ls='dashed')
plt.plot(listeThetaInc*180/np.pi,liste_r_TM,label="r_TM",color="red",ls='solid')
plt.plot(listeThetaInc*180/np.pi,liste_t_TM,label="t_TM",color="red",ls='dashed')
plt.xlabel('Angle incident $\\theta$ (°)')
plt.ylabel('Coefficient de Fresnel en amplitude')
plt.grid(1)
plt.legend()

# En intensité 
plt.figure()
plt.plot(listeThetaInc*180/np.pi,liste_R_TE,label="R_TE",color="blue",ls='solid')
plt.plot(listeThetaInc*180/np.pi,liste_T_TE,label="T_TE",color="blue",ls='dashed')
plt.plot(listeThetaInc*180/np.pi,liste_R_TM,label="R_TM",color="red",ls='solid')
plt.plot(listeThetaInc*180/np.pi,liste_T_TM,label="T_TM",color="red",ls='dashed')
plt.xlabel('Angle incident $\\theta$ (°)')
plt.ylabel('Coefficient de Fresnel en instensité')
plt.grid(1)
plt.legend()