import BeamPropagation as bp      # This is the script that handles the propagation
import sympy as sym               # For Symbolic examples
import numpy as np                # Handling of lists and for plotting
import matplotlib.pyplot as plt   # Plotting


"""A Gaussian beam can be defined by it's (radial) waist wo, it's Rayleigh range zR, and the location of its waist zO"""
w0 = 5.2E-3 # 1mm beam waist
lam = 532E-9 # wavelength of 355 nm (UV)
zR = bp.Zr(w0, lam) # Rayleigh range in m
z0 = 0 # location of waist in m

"""Define first 4-f optical system using matrices"""
d1, d2, d3, f1, f2 = sym.symbols('d1 d2 d3 f1 f2')
M = bp.mult(bp.prop(d3),bp.lens(f2),bp.prop(d2), bp.lens(f1), bp.prop(d1))

"""Use script to do all the ABCD and q-parameter math, and return the waist and radius of curvature functions"""
R, w = bp.q1_inv_func(0, w0, lam, M)

"""Substitute and extract the required separation between lenses of first 4-f system"""
distance_between_dmd_first_lens = 250E-3
first_focal_length = 250.9E-3
second_focal_length = 50E-3
demag = 1/5
target_w0 = demag*w0
w = w.subs(f1, first_focal_length).subs(f2, second_focal_length).subs(d1, distance_between_dmd_first_lens).subs(d3,0)
eq = sym.solve(w - target_w0, d2)[0]
distance_between_lens_of_first_4f = list(eq.atoms())[0]
print('Distance between lenses of first 4-f system = {} mm'.format(distance_between_lens_of_first_4f * 1E3))

# Sanity check
# expansion_factor = w.subs(d2,distance_between_lens_of_first_4f).subs(d3,0)/ w0
# print('beam is w = {:.2f} x w0'.format(expansion_factor))

# """Plot beam propagation up to 3 m after the first 4-f system"""
# M = bp.mult(bp.prop(d3),bp.lens(second_focal_length),bp.prop(distance_between_lens_of_first_4f), bp.lens(first_focal_length), bp.prop(distance_between_dmd_first_lens))
# R, w = bp.q1_inv_func(0, w0, lam, M)
# bp.plot(w,d3, rang = np.arange(0,0.050,.0005))


"""Define the full optical system of two 4-f setups using matrices"""
d1, d2, d3, d4, f1, f2, f3 = sym.symbols('d1 d2 d3 d4 f1 f2 f3')
M = bp.mult(bp.prop(d4),bp.lens(f3), bp.prop(d3),bp.lens(f2),bp.prop(d2), bp.lens(f1), bp.prop(d1))

"""Use script to do all the ABCD and q-parameter math, and return the waist and radius of curvature functions"""
R, w = bp.q1_inv_func(0, w0, lam, M)

# """Find the focal length of lens required after the first 4-f system to have a collimated beam, given a certain separation between the first 4-f system and this lens"""
distance_between_4fs = 550E-3
R_coll = R.subs(d1,distance_between_dmd_first_lens).subs(d2,distance_between_lens_of_first_4f).subs(d3,distance_between_4fs).subs(d4,0).subs(f1,first_focal_length).subs(f2,second_focal_length)
f3_coll = sym.solve(1/R_coll,f3)[0]
third_focal_length = list(f3_coll.atoms())[0]
print('For a fixed separation between first 4-f and third lens of {:.3f} mm, f3 = {:.3f} mm for a collimated beam'.format(distance_between_4fs* 1E3, third_focal_length * 1E3))

# # """Plot beam propagation up to 3 m after the first 4-f system"""
# M = bp.mult(bp.prop(d4),bp.lens(third_focal_length),bp.prop(distance_between_4fs), bp.lens(second_focal_length), bp.prop(distance_between_lens_of_first_4f), bp.lens(first_focal_length), bp.prop(distance_between_dmd_first_lens))
# R, w = bp.q1_inv_func(0, w0, lam, M)
# bp.plot(w,d4, rang = np.arange(0,0.050,.0005))

third_focal_length = 501.8E-3
R_coll = R.subs(d1,distance_between_dmd_first_lens).subs(d2,distance_between_lens_of_first_4f).subs(d4,0).subs(f1,first_focal_length).subs(f2,second_focal_length).subs(f3,third_focal_length)
d3_coll = sym.solve(1/R_coll,d3)[1]
distance_between_4fs = list(d3_coll.atoms())[0]
print('For a fixed third focal length of {:.3f} mm, d3 = {:.3f} mm, for a collimated beam'.format(third_focal_length* 1E3, distance_between_4fs * 1E3))

# """Plot beam propagation up to 3 m after the first 4-f system"""
# M = bp.mult(bp.prop(d4),bp.lens(third_focal_length),bp.prop(distance_between_4fs), bp.lens(second_focal_length), bp.prop(distance_between_lens_of_first_4f), bp.lens(first_focal_length), bp.prop(distance_between_dmd_first_lens))
# R, w = bp.q1_inv_func(0, w0, lam, M)
# bp.plot(w,d4, rang = np.arange(0,0.050,.0005))