import numpy as np
import matplotlib.pyplot as plt


def translation(coordEntree,n,dist):
    matTranslation=np.array([[1.0, dist/n],[0.0,1.0]])
    coordSortie=np.dot(matTranslation,coordEntree)
    return coordSortie

def ajoutTraj(traj,nouvCoord):
    traj=np.concatenate((traj,nouvCoord),axis=1)
    return traj

def dioptre_plan(coordEntree,n,AS,n_p,SAp):
    matTranslationSAp=np.array([[1.0, SAp/n_p],[0.0,1.0]])
    matTranslationAS=np.array([[1.0, AS/n],[0.0,1.0]])
    coordDioptre=translation(coordEntree,n,AS)
    matDioptrePlan=np.array([[1.0,0],[0,1.0]])
    matTransfer1=np.dot(matTranslationSAp,matDioptrePlan)
    matTransfer=np.dot(matTransfer1,matTranslationAS)
    coordSortie=np.dot(matTransfer,coordEntree)
    return coordSortie,coordDioptre

def dioptre_spherique(coordEntree,n,AS,n_p,SAp,SC):
    matTranslationSAp=np.array([[1.0, SAp/n_p],[0.0,1.0]])
    matTranslationAS=np.array([[1.0, AS/n],[0.0,1.0]])
    coordDioptre=translation(coordEntree,n,AS)
    phi=(n_p-n)/SC
    matDioptreSpherique=np.array([[1.0,0],[-phi,1.0]])
    matTransfer1=np.dot(matTranslationSAp,matDioptreSpherique)
    matTransfer=np.dot(matTransfer1,matTranslationAS)
    coordSortie=np.dot(matTransfer,coordEntree)
    return coordSortie,coordDioptre

def lentille_spherique(coordEntree,n,ASg,n_p,e,SAd,SCg,SCd):
    CoordMilieuLentille,coordDioptreG=dioptre_spherique(coordEntree,n,ASg,n_p,e/2,SCg)
    CoordSortie,coordDioptreD=dioptre_spherique(CoordMilieuLentille,n_p,e/2,n,SAd,SCd)
    return CoordSortie,coordDioptreD,coordDioptreG

def afficher_traj(z,traj):
    plt.plot(z,traj[0,:],color='blue')
    # plt.ylim(-2*max(traj[0,:]), 2*max(traj[0,:]))
    plt.ylim(-2, 2)
    plt.xlim(z[0],z[-1])
    plt.xlabel('z (m)')
    plt.ylabel('x (m)')
    plt.grid(1)
    
def desssinerAxeOptique(z):
    plt.plot([z[0],z[-1]],[0,0],color='black')
    plt.axvline(z[0],color='gray')
    plt.axvline(z[-1],color='gray')

def dessiner_dioptrePlan(position):
    plt.axvline(x=position,color='green')


def dessiner_dioptreSpherique(position,SC):
    if SC>0:
        theta = np.linspace(np.pi/2,3/2*np.pi, 100)
    else:
        theta = np.linspace(-np.pi/2,np.pi/2, 100)
    r=np.abs(SC)
    x1 = (position+r)+r*np.cos(theta)
    x2 = r*np.sin(theta)
    plt.plot(x1,x2,color='green')

def dessiner_lentille(positionCentre,SCg,SCd,e):
    if SCg>0:
        thetag = np.linspace(np.pi/2,3/2*np.pi, 100)
    else:
        thetag = np.linspace(-np.pi/2,np.pi/2, 100)
    r=np.abs(SCg)
    x1 = (positionCentre-e/2+SCg)+r*np.cos(thetag)
    x2 = r*np.sin(thetag)
    plt.plot(x1,x2,color='green')
    if SCd>0:
        thetad = np.linspace(np.pi/2,3/2*np.pi, 100)
    else:
        thetad = np.linspace(-np.pi/2,np.pi/2, 100)
    r=np.abs(SCd)
    x1 = (positionCentre+e/2+SCd)+r*np.cos(thetad)
    x2 = r*np.sin(thetag)
    plt.plot(x1,x2,color='green')
    plt.axvline(x=positionCentre,color='green',linestyle='--')
    
    

nMilieu=1.0
nVerre=1.5
ABEntree=1
phiEntree=-0.05
dist=20.0
z=[0,dist/2,dist]
coordEntree=np.array([[ABEntree],[nMilieu*phiEntree]])
repartitionABPlan =[1 ,0.75, 0.5, 0.25,0 ,-0.25, -0.5 ,-0.75 ,-1]



## Boucle pour rayon à travers dioptre plan
compteur=0;
for k in repartitionABPlan:
    Traj=0
    ABEntree=k;
    coordEntree=np.array([[ABEntree],[nMilieu*phiEntree]])
    coordSortieDioptre,coordDioptre=dioptre_plan(coordEntree,nMilieu,dist/2,nVerre,dist/2)
    Traj=ajoutTraj(coordEntree, coordDioptre)
    Traj=ajoutTraj(Traj, coordSortieDioptre)
    afficher_traj(z,Traj)
    desssinerAxeOptique(z)
    dessiner_dioptrePlan(dist/2)

## Boucle pour rayon à travers dioptre sphérique
compteur=0;
SC=5
plt.figure()
for k in repartitionABPlan:
    Traj=0
    ABEntree=k;
    coordEntree=np.array([[ABEntree],[nMilieu*phiEntree]])
    coordSortieDioptre,coordDioptre=dioptre_spherique(coordEntree,nMilieu,dist/2,nVerre,dist/2,SC)
    Traj=ajoutTraj(coordEntree, coordDioptre)
    Traj=ajoutTraj(Traj, coordSortieDioptre)
    afficher_traj(z,Traj)
    desssinerAxeOptique(z)
    dessiner_dioptreSpherique(dist/2,SC)
    
## Boucle pour rayon à travers lentille sphérique infini foyer
compteur=0;
SCg=5
SCd=-5
e=1
focale=1.0/((nVerre-nMilieu)*(1/SCg-1/SCd))
zLentille=[0, dist/2-e/2 ,dist/2+e/2 ,dist/2+focale]
plt.figure()
for k in repartitionABPlan:
    Traj=0
    ABEntree=k;
    coordEntree=np.array([[ABEntree],[nMilieu*phiEntree]])
    CoordSortie,coordDioptreD,coordDioptreG=lentille_spherique(coordEntree,nMilieu,dist/2-e/2,nVerre,e,focale-e/2,SCg,SCd)
    Traj=ajoutTraj(coordEntree, coordDioptreG)
    Traj=ajoutTraj(Traj, coordDioptreD)
    Traj=ajoutTraj(Traj, CoordSortie)
    afficher_traj(zLentille,Traj)
    desssinerAxeOptique(zLentille)
    dessiner_lentille(dist/2,SCg,SCd,e)
    
## Boucle pour rayon à travers lentille sphérique conjugaison
compteur=0;
OA=-30
repartitionPhiObj =np.linspace(-1/abs(OA), 1/abs(OA),9)
SCg=8
SCd=-8
e=1
focale=1.0/((nVerre-nMilieu)*(1/SCg-1/SCd))
inv_OA_p=1/focale+1/OA
OA_p=1/inv_OA_p
zLentille=[0, abs(OA)-e/2 ,abs(OA)+e/2 ,abs(OA)+OA_p]
plt.figure()
for k in repartitionPhiObj:
    Traj=0
    phiEntree=k;
    coordEntree=np.array([[0.0],[nMilieu*phiEntree]])
    CoordSortie,coordDioptreD,coordDioptreG=lentille_spherique(coordEntree,nMilieu,abs(OA)-e/2,nVerre,e,OA_p-e/2,SCg,SCd)
    Traj=ajoutTraj(coordEntree, coordDioptreG)
    Traj=ajoutTraj(Traj, coordDioptreD)
    Traj=ajoutTraj(Traj, CoordSortie)
    afficher_traj(zLentille,Traj)
    desssinerAxeOptique(zLentille)
    dessiner_lentille(abs(OA),SCg,SCd,e)