########################################################################
#
#    Christopher Miller
#    Laser Project
#
########################################################################

#import libraries
from pylab import *
from numpy import *

def plotlaser(G,p,f,b,plotnumber):
    clf()  #clears the current plot
    #The Laser 2D ODE
    def Laser_qdot(G,q,n,p,f,b):
        return G*n*q-b*q

    def Laser_ndot(G,q,n,p,f,b):
        return -G*n*q-f*n+p

    # 2D Fourth-Order Runge-Kutta Integrator
    def RKTwoD(G,q,n,fun1,fun2,p,f,b,dt):
        k1x = dt * fun1(G,q,n,p,f,b)
        k1y = dt * fun2(G,q,n,p,f,b)
        k2x = dt * fun1(G,q + k1x / 2.0,n + k1y / 2.0,p,f,b)
        k2y = dt * fun2(G,q + k1x / 2.0,n + k1y / 2.0,p,f,b)
        k3x = dt * fun1(G,q + k2x / 2.0,n + k2y / 2.0,p,f,b)
        k3y = dt * fun2(G,q + k2x / 2.0,n + k2y / 2.0,p,f,b)
        k4x = dt * fun1(G,q + k3x,n + k3y,p,f,b)
        k4y = dt * fun2(G,q + k3x,n + k3y,p,f,b)
        x = q + ( k1x + 2.0 * k2x + 2.0 * k3x + k4x ) / 6.0
        y = n + ( k1y + 2.0 * k2y + 2.0 * k3y + k4y ) / 6.0
        return x,y

    # Simulation parameters
    dt=.04

    # Set up arrays of iterates for several different initial conditions
    x1 = [0.4]
    y1 = [0.4]
    x2 = [0.4]
    y2 = [0.5]
    x3 = [0.4]
    y3 = [0.6]
    x4 = [0.5]
    y4 = [0.4]
    x5 = [0.5]
    y5 = [0.5]
    x6 = [0.5]
    y6 = [0.6]
    x7 = [0.6]
    y7 = [0.4]
    x8 = [0.6]
    y8 = [0.5]
    x9 = [0.6]
    y9 = [0.6]
    # Time
    t  = [ 0.0]
    # The number of time steps to integrate over
    N = 200


    # The main loop that generates the orbit, storing the states
    for n in xrange(0,N):
        # at each time step calculate new x(t) and y(t)
        x,y = RKTwoD(G,x1[n],y1[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x1.append(x)
        y1.append(y)
        x,y = RKTwoD(G,x2[n],y2[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x2.append(x)
        y2.append(y)
        x,y = RKTwoD(G,x3[n],y3[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x3.append(x)
        y3.append(y)
        x,y = RKTwoD(G,x4[n],y4[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x4.append(x)
        y4.append(y)
        x,y = RKTwoD(G,x5[n],y5[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x5.append(x)
        y5.append(y)
        x,y = RKTwoD(G,x6[n],y6[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x6.append(x)
        y6.append(y)
        x,y = RKTwoD(G,x7[n],y7[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x7.append(x)
        y7.append(y)
        x,y = RKTwoD(G,x8[n],y8[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x8.append(x)
        y8.append(y)
        x,y = RKTwoD(G,x9[n],y9[n],Laser_qdot,Laser_ndot,p,f,b,dt)
        x9.append(x)
        y9.append(y)
        t.append(t[n] + dt)

    if plotnumber == 1:    
        # Setup the parametric plot
        xlabel('q(t)') # set x-axis label
        ylabel('n(t)') # set y-axis label
        title('Laser Phase Chart between Cavity Photons (q) and Excited Atoms (n)')
        # Plot the trajectory in the phase plane
        plot(x1,y1,'b')
        plot(x2,y2,'r')
        plot(x3,y3,'g')
        plot(x4,y4,'b')
        plot(x5,y5,'r')
        plot(x6,y6,'g')
        plot(x7,y7,'b')
        plot(x8,y8,'r')
        plot(x9,y9,'g')
        #Display the plot in a window
        show()
    
    if plotnumber==2:
        xlabel('t') # set x-axis label
        ylabel('q(t)') # set y-axis label
        title('Laser Time Graph for Photon in Cavity (IC: q=.5,n=.5)')
        # Plot the trajectory in the phase plane
        plot(t,x5,'b')
        show()

    if plotnumber==3:
        q_range=arange(100)/10.0
        qdot=G*q_range*(.5-f*q_range+p)-b*q_range
        zerozero=0*q_range #used to graph a horizontal line

        xlabel('q(t)') # set x-axis label
        ylabel('q_dot(t)') # set y-axis label
        title('Laser Stability Diagram for Photons in the Cavity (Nmax=.5)')
        # Plot the trajectory in the phase plane
        plot(q_range,qdot,'b')
        plot(q_range,zerozero,'g')
        axis([0, 2, -.5, .5])
        show()