from math import *
from pylab import *
from IndiceVerre import *

class Rayon:
    def __init__(self,P,U):
        self.P = P
        N = sqrt(U[0]*U[0]+U[1]*U[1]+U[2]*U[2])
        self.U = [U[0]/N,U[1]/N,U[2]/N]
        self.actif = True
        self.longueur = 0.0
    def getPoint(self,t):
        return [self.P[0]+t*self.U[0],self.P[1]+t*self.U[1],self.P[2]+t*self.U[2]]
                
class RayonComplet:
    def __init__(self):
        self.listeR = []
        self.n = 0
    def ajouter(self,R):
        self.listeR.append(R)
        self.n += 1
    def getRayon(self):
        if self.n==0:
            return 0
        else:
            return self.listeR[self.n-1]
            
    def getXYZ(self,tf):
        x=[]
        y=[]
        z=[]
        for k in range(self.n):
           x.append(self.listeR[k].P[0])
           y.append(self.listeR[k].P[1])
           z.append(self.listeR[k].P[2])
        if self.listeR[self.n-1].actif:
            P = self.listeR[self.n-1].getPoint(tf)
            x.append(P[0])
            y.append(P[1])
            z.append(P[2])
        return [x,y,z] 
                
    def longueurOptique(self,last):
        l = 0.0
        for k in range(self.n-1-last):
            l += self.listeR[k].longueur
        return l
                
class Faisceau:
    def __init__(self):
        self.listeRayons = []
        self.n = 0
    def ajouter(self,Rc):
        self.listeRayons.append(Rc)
        self.n += 1
                
    def barycentre(self):
        x = 0
        y = 0
        na = 0
        for Rc in self.listeRayons:
            R = Rc.getRayon()
            if R.actif:
                na +=1 
                x += R.P[0]
                y += R.P[1]
        return [x/na,y/na]
                
class FaisceauParalleleMeridien(Faisceau):
    def __init__(self,nr,alpha,z,r):
        Faisceau.__init__(self)
        dr = 2.0*r/(nr-1)
        cosinus = cos(alpha)
        sinus = sin(alpha)
        for k in range(nr):
            r1 = -r+k*dr
            x = r1*cosinus
            z = -r1*sinus + z
            Rc = RayonComplet()
            Rc.ajouter(Rayon([x,0,z],[sinus,0,cosinus]))
            self.ajouter(Rc)
                
class FaisceauParalleleCylindre(Faisceau):
    def __init__(self,nr,alpha,z,r):
        Faisceau.__init__(self)
        N = int(nr/2)
        cosinus = cos(alpha)
        sinus = sin(alpha)
        dtheta = pi/(N-1)
        for k in range(2*N):
            theta = dtheta*k
            xx = r*cos(theta)
            yy = r*sin(theta)
            xp = xx*cosinus
            yp = yy
            zp = -xx*sinus + z
            Rc = RayonComplet()
            Rc.ajouter(Rayon([xp,yp,zp],[sinus,0,cosinus]))
            self.ajouter(Rc)
        
class FaisceauParallele(Faisceau):
    def __init__(self,nr,alpha,z,r):
        Faisceau.__init__(self)
        dx = 2*r/nr
        r2 = r*r
        cosinus = cos(alpha)
        sinus = sin(alpha)
        for i in range(nr):
            xx = -r+dx*i
            for j in range(nr):
                yy = -r+dx*j
                rr2 = xx*xx+yy*yy
                if (rr2 < r2):
                    xp = xx*cosinus
                    yp = yy
                    zp = -xx*sinus + z
                    Rc = RayonComplet()
                    Rc.ajouter(Rayon([xp,yp,zp],[sinus,0,cosinus]))
                    self.ajouter(Rc)
        
class FaisceauDivergentMeridien(Faisceau):
    def __init__(self,nr,zo,xo,z,r):
        dr = 2*r/(nr-1)
        for k in range(nr):
            r1 = -r+k*dr
            Rc = RayonComplet()
            if z > zo:
                Rc.ajouter(Rayon([xo,0,z0],[r1-xo,0,z-zo]))
            else:
                u = xo-r1
                w = zo-z
                f = 2*w/sqrt(u*u+w*w)
                Rc.ajouter(Rayon([xo-u*f,0,zo-w*f],[u,0,w]))
            self.ajouter(Rc)
            
class FaisceauDivergent(Faisceau):
    def __init__(self,nr,zo,xo,z,r):
        dx = 2*r/nr;
        r2 = r*r
        for i in range(nr):
            xx = -r+dx*i
            for j in range(nr):
                yy = -r+dx*j
                rr2 = xx*xx+yy*yy
                if rr2 < r2:
                    Rc = RayonComplet()
                    if z > zo:
                        Rc.ajouter(Rayon([xo,0,z0],[xx-xo,yy,z-zo]))
                    else:
                        u = xo-xx
                        v = -yy
                        w = zo-z
                        f = 2*w/sqrt(u*u+v*v+w*w)
                        Rc.ajouter(Rayon([xo-u*f,-v*f,zo-w*f],[u,v,w]))
                self.ajouter(Rc)
            
class Dioptre:
    def __init__(self,z,c,n1,n2,r):
        self.z = z
        self.c = c
        self.n1 = n1
        self.n2 = n2
        self.r = r
        self.r2 = r*r
                
    def matriceRefraction(self,raie):
        return [[1,self.c*(self.n1.indice(raie)-self.n2.indice(raie))],[0,1]]
    def matriceTranslation(self,z,raie):
        return [[1,0],[(self.z-z)/self.n1.indice(raie),1]]
            
    def dzBord(self):
        if self.c==0:
            return 0.0
        rc = abs(1/self.c)
        dz = rc-sqrt(rc*rc-self.r*self.r)
        if self.c < 0:
            dz = -dz
        return dz
            
class Spherique(Dioptre):
    def refraction(self,R1,raie):
        n1 = self.n1.indice(raie)
        n2 = self.n2.indice(raie)
        z1s = R1.P[2]-self.z
        R1.actif = False
        if self.c==0:
            if (R1.U[2]!=0):
                t = -z1s/R1.U[2]
                P2 = R1.getPoint(t)
            else:
                return R1
        else:
            B = (R1.P[0]*R1.U[0]+R1.P[1]*R1.U[1]+z1s*R1.U[2])-R1.U[2]/self.c
            D = (R1.P[0]*R1.P[0]+R1.P[1]*R1.P[1]+z1s*z1s)-2*z1s/self.c
            delta = B*B-D
            if (delta <=0):
                return R1
            sqd = sqrt(delta)
            ta = -B+sqd
            tb = -B-sqd
            zc = self.z+1.0/self.c
            P2 = R1.getPoint(ta)
            a = (P2[2]-zc)*(self.z-zc)
            P2 = R1.getPoint(tb)
            b = (P2[2]-zc)*(self.z-zc)
            if ((a<0)|(ta<0)):
                t = tb
                P2 = R1.getPoint(t)
            elif ((b<0)|(tb<0)):
                t = ta
                P2 = R1.getPoint(t)
            else:
                if (ta < tb):
                    t = ta
                    P2 = R1.getPoint(t)
                else:
                    t = tb
                    P2 = R1.getPoint(t)
        R1.longueur = t*abs(n1)
        if (P2[0]*P2[0]+P2[1]*P2[1] > self.r2):
            return R1
        cosTheta1 = R1.U[2]*(1+self.c*(self.z-P2[2]))-self.c*(R1.U[0]*P2[0]+R1.U[1]*P2[1])
        x = 1-pow(n1/n2,2)*(1-cosTheta1*cosTheta1)
        if x<0:
            return R1
        cosTheta2 = sqrt(x)
        if R1.U[2]>0:
            K = n2*cosTheta2-n1*cosTheta1
        else:
            cosTheta1 = -cosTheta1
            K = -n2*cosTheta2+n1*cosTheta1
        U2 = [0,0,0]
        U2[0] = (n1*R1.U[0]-self.c*K*P2[0])/n2
        U2[1] = (n1*R1.U[1]-self.c*K*P2[1])/n2
        U2[2] = (n1*R1.U[2]-self.c*K*(P2[2]-self.z)+K)/n2
        if n2*n1<0:
            U2[0] = -U2[0]
            U2[1] = -U2[1]
            U2[2] = -U2[2]
        R1.actif = True
        return Rayon(P2,U2)
            
    def profil(self,npts):
        x = []
        z = []
        if self.c==0:
            x = [-self.r,self.r]
            z = [self.z,self.z]
        else:
            Rc = 1.0/self.c
            zc = self.z + Rc
            alpha = asin(fabs(self.r/Rc))
            if Rc<0:
                alpha1 = -alpha
                alpha2 = alpha
            else:
                alpha1 = pi-alpha
                alpha2 = pi+alpha
            da = (alpha2-alpha1)/(npts-1)
            Rc = fabs(Rc)
            for k in range(npts):
                a = alpha1 + da*k
                z.append(zc+Rc*cos(a))
                x.append(Rc*sin(a))
        return [x,z]
            
class Conique(Dioptre):
    def __init__(self,z,c,S,n1,n2,r):
        self.S = S
        Dioptre.__init__(self,z,c,n1,n2,r)
    def refraction(self,R1,raie):
        n1 = self.n1.indice(raie)
        n2 = self.n2.indice(raie)
        R1.actif = False
        z1s = R1.P[2]-self.z
        A = self.c/2.0*(R1.U[0]*R1.U[0]+R1.U[1]*R1.U[1]+self.S*R1.U[2]*R1.U[2])
        B = self.c*(R1.P[0]*R1.U[0]+R1.P[1]*R1.U[1]+self.S*z1s*R1.U[2])-R1.U[2]
        D = self.c/2.0*(R1.P[0]*R1.P[0]+R1.P[1]*R1.P[1]+self.S*z1s*z1s)-z1s
        if A==0.0:
            if B==0.0:
                return R1
            t = -D/B
            P2 = R1.getPoint(t)
        else:
            delta = B*B-4*A*D
            if delta<=0.0:
                return R1
            sqd = sqrt(delta)
            ta = (-B+sqd)/(2*A)
            tb = (-B-sqd)/(2*A)
            Pa = R1.getPoint(ta)
            Pb = R1.getPoint(tb)
            if self.S==0.0:
                a=1
                b=1
            else:
                zc = z-1.0/(self.c*self.S)
                a = (Pa[2]-zc)*(self.z-zc)
                b = (Pb[2]-zc)*(self.z-zc)
            if ((a<0)|(ta<0)):
                t = tb
                P2 = R1.getPoint(t)
            elif ((b<0)|(tb<0)):
                t = ta
                P2 = R1.getPoint(t)
            else:
                if (ta < tb):
                    t = ta
                    P2 = R1.getPoint(t)
                else:
                    t = tb
                    P2 = R1.getPoint(t)
        R1.longueur = t*abs(n1)
        r2 = P2[0]*P2[0]+P2[1]*P2[1]
        if ( r2 > self.r2):
            return R1
        az = 1-self.c*self.S*(self.z-P2[2])
        den = sqrt(self.c*self.c*r2+az*az)
        cosTheta1 = (R1.U[2]*az-self.c*(R1.U[0]*P2[0]+R1.U[1]*P2[1]))/den
        x = 1-pow(n1/n2,2)*(1-cosTheta1*cosTheta1)
        if x<0:
            return R1
        cosTheta2 = sqrt(x)
        if R1.U[2]>0:
            K = n2*cosTheta2-n1*cosTheta1
        else:
            cosTheta1 = -cosTheta1
            K = -n2*cosTheta2+n1*cosTheta1
        U2 = [0,0,0]
        U2[0] = (n1*R1.U[0]-self.c*K/den*P2[0])/n2
        U2[1] = (n1*R1.U[1]-self.c*K/den*P2[1])/n2
        U2[2] = (n1*R1.U[2]+K/den*az)/n2
        if n2*n1<0:
            U2[0] = -U2[0]
            U2[1] = -U2[1]
            U2[2] = -U2[2]
        R1.actif = True
        return Rayon(P2,U2)
                
    def profil(self,npts):
        x = []
        z = []
        if self.c==0:
            x = [-self.r,self.r]
            z = [self.z,self.z]
        else:
            Rc = 1.0/self.c
            zc = self.z + Rc
            alpha = asin(fabs(self.r/Rc))
            if Rc<0:
                alpha1 = -alpha
                alpha2 = alpha
            else:
                alpha1 = pi-alpha
                alpha2 = pi+alpha
            da = (alpha2-alpha1)/(npts-1)
            Rc = fabs(Rc)
            for k in range(npts):
                a = alpha1 + da*k
                z.append(zc+Rc*cos(a))
                x.append(Rc*sin(a))
        return [x,z]
            
class Ecran(Dioptre):
    def __init__(self,z,r):
        Dioptre.__init__(self,z,0,IndiceVide(),IndiceVide(),r)
    def refraction(self,R1,raie):
        t = (self.z-R1.P[2])/R1.U[2]
        if (t<0):
            R1.actif = False
            return R1
        na = self.n1.indice(raie)
        R1.longueur = t*abs(na)
        P2 = R1.getPoint(t)
        U2 = R1.U
        R2 = Rayon(P2,U2)
        if (P2[0]*P2[0]+P2[1]*P2[1] < self.r2):
            R2.actif = False
        return R2
    def profil(self,npts):
        x = [-self.r,self.r]
        z = [self.z,self.z]
        return [x,z]
                
class PlanImage(Dioptre):
    def __init__(self,z,r):
        Dioptre.__init__(self,z,0,IndiceVide(),IndiceVide(),r)
    def refraction(self,R1,raie):
        t = (self.z-R1.P[2])/R1.U[2]
        na = self.n1.indice(raie)
        R1.longueur = t*abs(na)
        P2 = R1.getPoint(t)
        U2 = R1.U
        R2 = Rayon(P2,U2)
        return R2
    def profil(self,npts):
        x = [-self.r,self.r]
        z = [self.z,self.z]
        return [x,z]
                
class Diaphragme(Dioptre):
    def __init__(self,z,rmax):
        self.rmax = rmax
        Dioptre.__init__(self,z,0,IndiceVide(),IndiceVide(),r)
    def refraction(self,R1,raie):
        t = (self.z-R1.P[2])/self.U[2]
        if (t<0):
            R1.actif = False
            return R1
        na = self.n1.indice(raie)
        R1.longueur = t*abs(na)
        P2 = R1.getPoint(t)
        U2 = R1.U
        R2 = Rayon(P2,U2)
        if (P2[0]*P2[0]+P2[1]*P2[1] > self.r2):
            R2.actif = False
        return R2
    def profil(self,npts):
        x = [[self.r,self.rmax],[-self.r,-self.rmax]]
        z = [[self.z,self.z],[self.z,self.z]]
        return [x,z]
                
class SphereImage(Dioptre):
    def __init__(self,z,PC):
        self.PC = PC
        Dioptre.__init__(self,z,0,IndiceVide(),IndiceVide(),0)
    def refraction(self,R1,raie):
        b = R1.U[0]*(R1.P[0]-self.PC[0])+R1.U[1]*(R1.P[1]-self.PC[1])+R1.U[2]*(R1.P[2]-self.PC[2])
        a0 = R1.P[0]-self.PC[0]
        a1 = R1.P[1]-self.PC[1]
        a2 = R1.P[2]-self.PC[2]
        R = self.PC[2]-self.z
        d = a0*a0+a1*a1+a2*a2-R*R
        delta = b*b-d
        if delta <= 0:
            R1.actif = False
            return R1
        rdelta = sqrt(delta)
        t = min(-b+rdelta,-b-rdelta)
        if t < 0:
            R1.actif = False
            return R1
        na = self.n1.indice(raie)
        R1.longueur = t*abs(na)
        P2 = R1.getPoint(t)
        U2 = R1.U
        R2 = Rayon(P2,U2)
        return R2
    def profil(self,npts):
        return [[0],[0]]
                
class Position:
    def __init__(self,z):
        self.z = z
        self.infini = False
                
class SystemeCentre:
    def __init__(self):
        self.listeD = []
        self.nbD = 0
        self.M = [[1,0],[0,1]]
        self.no = Vide()
        self.ni = Vide()
    def ajouter(self,dioptre):
        self.listeD.append(dioptre)
        self.nbD += 1
    def ajouterListe(self,liste):
        for d in liste:
            self.ajouter(d)
    def refraction(self,rayonComplet,raie):
        k = 0
        R1 = rayonComplet.getRayon()
        actif = R1.actif
        while k < self.nbD and actif:
            R2 = self.listeD[k].refraction(R1,raie)
            rayonComplet.ajouter(R2)
            R1 = R2
            k += 1
            actif = R2.actif
    def refractionFaisceau(self,faisceau,raie):
        for Rc in faisceau.listeRayons:
            self.refraction(Rc,raie)
                
    def matriceTransfert(self,raie):
        self.M = self.listeD[0].matriceRefraction(raie)
        zd = self.listeD[0].z
        k = 1
        while k < self.nbD:
            MT = self.listeD[k].matriceTranslation(zd,raie)
            self.M = self.multiplier(MT,self.M)
            MR = self.listeD[k].matriceRefraction(raie)
            self.M = self.multiplier(MR,self.M)
            zd = self.listeD[k].z
            k += 1
        return self.M
    def multiplier(self,A,B):
        C = [[0,0],[0,0]]
        C[0][0] = A[0][0]*B[0][0]+A[0][1]*B[1][0]
        C[0][1] = A[0][0]*B[0][1]+A[0][1]*B[1][1]
        C[1][0] = A[1][0]*B[0][0]+A[1][1]*B[1][0]
        C[1][1] = A[1][0]*B[0][1]+A[1][1]*B[1][1]
        return C
                
    def image(self,PObjet,raie):
        PImage = Position(0)
        if PObjet.infini:
            if self.M[0][1]==0:
                PImage.infini = True
            else:
                ti = -self.ni.indice(raie)*self.M[1][1]/self.M[0][1]
                PImage.z = self.listeD[self.nbD-1].z+ti
        else:
            to = PObjet.z-self.listeD[0].z
            den = self.M[0][0]-self.M[0][1]*to/self.no.indice(raie)
            if den==0:
                PImage.infini = True
            else:
                ti = self.ni.indice(raie)*(-self.M[1][0]+self.M[1][1]*to/self.no.indice(raie))/den
                PImage.z = self.listeD[self.nbD-1].z+ti
        return PImage
                
                
    def grandissement(self,PObjet,raie):
        if PObjet.infini:
            if self.M[0][1]==0:
                return self.M[0][0]*self.no.indice(raie)/self.ni.indice(raie)
            return (-self.M[1][1]/self.M[0][1]*self.M[0][0]+self.M[1][0])*self.no.indice(raie)
        else:
            to = PObjet.z-self.listeD[0].z
            den = self.M[0][0]-self.M[0][1]*to/self.no.indice(raie)
            if den==0:
                return self.M[0][1]/self.ni.indice(raie)
            else:
                return 1.0/den
                
    def foyers(self,raie):
        Fo = Position(0)
        Fi = Position(0)
        Ho = Position(0)
        Hi = Position(0)
        if self.M[0][1]==0:
            Fo.infini = True
            Fi.infini = True
            Ho.infini = True
            Hi.infini = True
        else:
            Fo.z = self.listeD[0].z + self.no.indice(raie)*self.M[0][0]/self.M[0][1]
            Fi.z = self.listeD[self.nbD-1].z - self.ni.indice(raie)*self.M[1][1]/self.M[0][1]
            Ho.z = self.listeD[0].z + self.no.indice(raie)*(self.M[0][0]-1)/self.M[0][1]
            Hi.z = self.listeD[self.nbD-1].z + self.ni.indice(raie)*(1-self.M[1][1])/self.M[0][1]
        return [Fo,Fi,Ho,Hi]
                
    def tracerProfils(self,npts,color):
        for d in self.listeD:
            [x,z] = d.profil(npts)
            plot(z,x,color='k')
                
    def pointConvergence(self,z0,r,raie):
        if r==0.0:
            self.matriceTransfert(raie)
            [Fo,Fi,Ho,Hi] = self.foyers(raie)
            return Fi.z
        else:
            R = Rayon([r,0,z0],[0,0,1.0])
            Rc = RayonComplet()
            Rc.ajouter(R)
            self.refraction(Rc,raie)
            R2 = Rc.getRayon()
            t = -R2.P[0]/R2.U[0]
            z = R2.P[2]+R2.U[2]*t
            return z