red_car_mass   := 251.8~0.05 "CAR 1"
blue_car_mass  := 254.0~0.05 "CAR 2"
m1_mass := 252.4~0.05
m2_mass  := 254.2~0.05

slopes := [0.0186~0.00015, 0.0213~0.00025, 0.0259~0.00024, 0.0187~0.00069, 0.0225~0.00025]

define extract_column(matrix, col) as
    column := []

    for row in matrix loop
        column | row[col]
    end loop

    return column
end extract_column

define extract_columns(matrices, col) as
    return ((x) -> extract_column(x, col))`(matrices)
end extract_columns

define flatten(matrix) as
    result := []

    for set in matrix loop
        result.extend(set)
    end loop

    return result
end flatten

define sum(x, y) as
    return x + y
end sum

define sums(x, y) as
    return sum`(x, y)
end sums


print("The standard deviation of our measurements is:", std(slopes))


"=== Elastic data ==="

"[beginning_pos, cart2_v_at, cart1_v_after, cart2_v_after]"
elastic_collisions_data_massless := [
    [1.6,   0.312249, -0.289780, -0.005616],
    [2.070, 0.343593, -0.323470, -0.002712],
    [1.94,  0.365498, -0.313415, -0.001286],
    [1.26,  0.688140, -0.708082, 0.008666],
    [2.04,  0.205810, -0.186434, -0.011421]
]

"[beginning_pos, cart2_v_at, cart1_v_after, cart2_v_after]"
elastic_collisions_data_massed := [
    [1.3, 0.338302,   -0.330059, -0.003630],
    [0.860, 0.383199, -0.34254,  0.001001],
    [1.260, 0.260496, -0.256836, -0.006353],
    [1.270, 0.472937, -0.460785, -0.002867],
    [1.260, 0.574692, -0.559910, 0.0001650842]
]

"[beginning_pos, cart2_v_at, cart1_v_after, cart2_v_after]"
elastic_collisions_data_c1m1 := [
    [1.26, 0.447398, -0.284034, -0.14181],
    [1.08, 0.443326, -0.272659, -0.134786],
    [1.34, 0.435538, -0.266906, -0.146262],
    [1.4,  0.376664, -0.244017, -0.118208]
]

"[beginning_pos, cart2_v_at, cart1_v_after, cart2_v_after]"
elastic_collisions_data_c1m2 := [
    [1.7,  0.596225, -0.289686, -0.300989],
    [0.7,  0.59995,  -0.278186, -0.267378],
    [0.68, 0.785811, -0.370249, -0.351386],
    [1.72, 0.492266, -0.220981, -0.233584]
]

"[beginning_pos, cart2_v_at, cart1_v_after, cart2_v_after]"
elastic_collisions_data_c2m1 := [
    [0.71,  0.302166, -0.41089, 0.089709],
    [1.640, 0.220673, -0.300117, 0.061775],
    [0.880, 0.307913, -0.418068, 0.100391],
    [0.580, 0.35472,  -0.446097, 0.117493]
]

"[beginning_pos, cart2_v_at, cart1_v_after, cart2_v_after]"
elastic_collisions_data_c2m2 := [
    [1.32, 0.197016, -0.268901, 0.083176],
    [1.19, 0.26225,  -0.321149, 0.111764],
    [1.46, 0.229498, -0.303468, 0.094574],
    [1.33, 0.203483, -0.286996, 0.086106]
]


"=== Inelastic data ==="

"[beginning_pos, cart2_v_at, cart1_v_after]"
inelastic_collisions_data_massless := [
    [1.43, 0.392352, -0.196709],
    [1.35, 0.496739, -0.241296],
    [1.08, 0.69781,  -0.337878],
    [1.03, 0.489847, -0.239846]
]

"[beginning_pos, cart2_v_at, cart1_v_after]"
inelastic_collisions_data_c1m1 := [
    [0.960, 0.588607, -0.202218],
    [1.88,  0.407084, -0.14504],
    [1.14,  0.580273, -0.189352]
]

"[beginning_pos, cart2_v_at, cart1_v_after]"
inelastic_collisions_data_c1m2 := [
    [0.96, 0.851525, -0.132017],
    [2.28, 0.558863, -0.136012],
    [2.1,  0.632977, -0.153357]
]

"[beginning_pos, cart2_v_at, cart1_v_after]"
inelastic_collisions_data_c2m1 := [
    [2.08, 0.442954, -0.294568],
    [2.56, 0.47818,  -0.306290],
    [1.26, 0.460557, -0.315441]
]

"[beginning_pos, cart2_v_at, cart1_v_after]"
inelastic_collisions_data_c2m2 := [
    [1.84, 0.56718,  -0.393595],
    [1.78, 0.413264, -0.31968],
    [1.74, 0.492195, -0.38198]
]

all_data := [
    elastic_collisions_data_massless,
    elastic_collisions_data_c1m1,
    elastic_collisions_data_c1m2,
    elastic_collisions_data_c2m1,
    elastic_collisions_data_c2m2,
    inelastic_collisions_data_massless,
    inelastic_collisions_data_c1m1,
    inelastic_collisions_data_c1m2,
    inelastic_collisions_data_c2m1,
    inelastic_collisions_data_c2m2
]


define calc_kinetic(mass, velocity) as
    return 0.5 * mass * velocity ^ 2
end calc_kinetic

define calc_momentum(mass, velocity) as
    return mass * velocity
end calc_momentum


define negate_car1v(list) as
    for row in list loop
        row[2] := -1 * row[2]
    end loop
end negate_car1v


"Since the sensors were in different directions, we needed to negate one of them"
"We chose to negate the one pointing towards car 1"
negate_car1v`(all_data)

from physics_utils.graph import SimpleGraph, show

"From the velocities of the carts, calculate the momenta and kinetic energies. Plot the total kinetic energy before the"
"collision vs the total kinetic energy after collision. Is kinetic energy conserved (i.e. is the collision 'elastic')?"
"What slope would this correspond to? Also, plot the momenta before collision vs momenta after the"
"collision. Is momentum conserved?"

elastic_datasets := [
    elastic_collisions_data_massless,
    elastic_collisions_data_c1m1,
    elastic_collisions_data_c1m2,
    elastic_collisions_data_c2m1,
    elastic_collisions_data_c2m2
]

elastic_masses_c1 := [
    red_car_mass,
    red_car_mass + m1_mass,
    red_car_mass + m1_mass + m2_mass,
    red_car_mass,
    red_car_mass
]

elastic_masses_c2 := [
    blue_car_mass,
    blue_car_mass,
    blue_car_mass,
    blue_car_mass + m1_mass,
    blue_car_mass + m1_mass + m2_mass
]



elastic_c1_velocities_after  := extract_columns(elastic_datasets, 2)
elastic_c2_velocities_before := extract_columns(elastic_datasets, 1)
elastic_c2_velocities_after  := extract_columns(elastic_datasets, 3)

"First, let's calculate the momenta for each dataset:"

    print()

    define calc_momenta_on_dataset(mass, velocities) as
        return ((x) -> calc_momentum(mass, x))`(velocities)
    end calc_momenta_on_dataset

    cart1_momenta_after  := calc_momentum**(*elastic_masses_c1, **elastic_c1_velocities_after) 
    cart2_momenta_before := calc_momentum**(*elastic_masses_c2, **elastic_c2_velocities_before)
    cart2_momenta_after  := calc_momentum**(*elastic_masses_c2, **elastic_c2_velocities_after)

    total_momenta_before := cart2_momenta_before
    total_momenta_after  := sums`(cart2_momenta_after, cart1_momenta_after)
    elastic_total_momenta_before := total_momenta_before
    elastic_total_momenta_after := total_momenta_after

    print("Total momenta before:", total_momenta_before)
    print("Total momenta after:", total_momenta_after)

"& graph it:"

    momenta_graph := SimpleGraph("Momentum before vs after elastic collision")
    momenta_graph.set_x_axis(flatten(total_momenta_after), "Total momentum after (g m / s)")
    momenta_graph.set_y_axis(flatten(total_momenta_before), "Total momentum before (g m / s)")

    momenta_graph.plot_points()
    momenta_graph.put_labels()

    line_info := momenta_graph.best_fit(true)

    print("Line of best fit for elastic momenta graph:", "$y=({})x+{}$".format(line_info[0].latex(false), line_info[1].latex(false)))

"Now, let's calculate the kinetic energy:"

    print()

    define calc_ke_on_dataset(mass, velocities) as
        return ((x) -> calc_kinetic(mass, x))`(velocities)
    end calc_ke_on_dataset

    cart1_ke_after  := calc_ke_on_dataset`(elastic_masses_c1, elastic_c1_velocities_after)
    cart2_ke_before := calc_ke_on_dataset`(elastic_masses_c2, elastic_c2_velocities_before)
    cart2_ke_after  := calc_ke_on_dataset`(elastic_masses_c2, elastic_c2_velocities_after)

    total_ke_before := cart2_ke_before
    total_ke_after  := sums`(cart2_ke_after, cart1_ke_after)
    elastic_total_ke_before := total_ke_before
    elastic_total_ke_after := total_ke_after

    print("Total kinetic energy before:", total_ke_before)
    print("Total kinetic energy after:", total_ke_after)

"& graph it:"

    ke_graph := SimpleGraph("Kinetic energy before vs after elastic collision")
    ke_graph.set_x_axis(flatten(total_ke_after), "Total kinetic energy after (g m^2 / s^2)")
    ke_graph.set_y_axis(flatten(total_ke_before), "Total kinetic energy before (g m^2 / s^2)")

    ke_graph.plot_points()
    ke_graph.put_labels()

    line_info := ke_graph.best_fit(true)

    print("Line of best fit for elastic kinetic graph:", "$y=({})x+{}$".format(line_info[0].latex(false), line_info[1].latex(false)))

"Ok now we're doing the stuff for the inelastic collision"

"you will find each cart's velocity before the collision, and the velocity of the"
"combined system of two carts after collision (because they're stuck together). Once again, compare"
"the momenta and kinetic energies before and after the collision. Is this collision elastic? Is kinetic"
"energy conserved? Is momentum conserved? If one was not conserved, what happened to it?"

inelastic_datasets := [
    inelastic_collisions_data_massless,
    inelastic_collisions_data_c1m1,
    inelastic_collisions_data_c1m2,
    inelastic_collisions_data_c2m1,
    inelastic_collisions_data_c2m2
]

inelastic_masses_c1 := [
    red_car_mass,
    red_car_mass + m1_mass,
    red_car_mass + m1_mass + m2_mass,
    red_car_mass,
    red_car_mass
]

inelastic_masses_c2 := [
    blue_car_mass,
    blue_car_mass,
    blue_car_mass,
    blue_car_mass + m1_mass,
    blue_car_mass + m1_mass + m2_mass
]

inelastic_c1_velocities_after  := extract_columns(inelastic_datasets, 2)
inelastic_c2_velocities_before := extract_columns(inelastic_datasets, 1)
inelastic_c2_velocities_after  := extract_columns(inelastic_datasets, 2)

"Let's calculate the momenta:"

    print()

    total_momenta_before := calc_momenta_on_dataset`(inelastic_masses_c2, inelastic_c2_velocities_before)
    total_momenta_after  := calc_momenta_on_dataset`(sums(inelastic_masses_c2, inelastic_masses_c1), inelastic_c2_velocities_after)
    inelastic_total_momenta_before := total_momenta_before
    inelastic_total_momenta_after := total_momenta_after

    print("Total momenta before:", total_momenta_before)
    print("Total momenta after:", total_momenta_after)

"& graph it:"

    momenta_graph := SimpleGraph("Momentum before vs after inelastic collision")
    momenta_graph.set_x_axis(flatten(total_momenta_after), "Total momentum after (g m / s)")
    momenta_graph.set_y_axis(flatten(total_momenta_before), "Total momentum before (g m / s)")

    momenta_graph.plot_points()
    momenta_graph.put_labels()

    line_info := momenta_graph.best_fit(true)

    print("Line of best fit for inelastic momentum graph:", "$y=({})x+{}$".format(line_info[0].latex(false), line_info[1].latex(false)))

"Now, let's calculate the kinetic energy:"

    print()

    cart1_ke_after  := calc_ke_on_dataset`(inelastic_masses_c1, inelastic_c1_velocities_after)
    cart2_ke_before := calc_ke_on_dataset`(inelastic_masses_c2, inelastic_c2_velocities_before)
    cart2_ke_after  := calc_ke_on_dataset`(inelastic_masses_c2, inelastic_c2_velocities_after)

    total_ke_before := cart2_ke_before
    total_ke_after  := sums`(cart2_ke_after, cart1_ke_after)
    inelastic_total_ke_before := total_ke_before
    inelastic_total_ke_after := total_ke_after

    print("Total kinetic energy before:", total_ke_before)
    print("Total kinetic energy after:", total_ke_after)

"& graph it:"

    ke_graph := SimpleGraph("Kinetic energy before vs after inelastic collision")
    ke_graph.set_x_axis(flatten(total_ke_after), "Total kinetic energy after (g m^2 / s^2)")
    ke_graph.set_y_axis(flatten(total_ke_before), "Total kinetic energy before (g m^2 / s^2)")

    ke_graph.plot_points()
    ke_graph.put_labels()

"Finally, let's see our graphs:"
show()

print()
print("Now, for our tables")
print()

from physics_utils.table import Table2D

trial_num := [0]

define make_row(m1, m2, run) as
    trial_num[0] := trial_num[0] + 1
    return [trial_num[0].__int__(), m1, m2, run[1], run[2], run[3]]
end make_row

table_data := flatten(((m1, m2, runs) -> ((run) -> make_row(m1, m2, run))`(runs))`(elastic_masses_c1, elastic_masses_c2, elastic_datasets))

elastic_collision_table := Table2D()

elastic_collision_table.set_labels(["$m_1$", "$m_2$", "$v_{i_2}$", "$v_{f_1}$", "$v_{f_2}$"])
elastic_collision_table.set_data(table_data)

"print(elastic_collision_table.latex())"


define make_row(m1, m2, run) as
    trial_num[0] := trial_num[0] + 1
    return [trial_num[0].__int__(), m1, m2, run[1], run[2]]
end make_row



table_data := flatten(((m1, m2, runs) -> ((run) -> make_row(m1, m2, run))`(runs))`(inelastic_masses_c1, inelastic_masses_c2, inelastic_datasets))

inelastic_collision_table := Table2D()

inelastic_collision_table.set_labels(["$m_1$", "$m_2$", "$v_{i_2}$", "$v_{f_1}$"])
inelastic_collision_table.set_data(table_data)

"print(inelastic_collision_table.latex())"

trial_num[0] := 0

define make_row(momenta_before, momenta_after, ke_before, ke_after) as
    trial_num[0] := trial_num[0] + 1
    return [trial_num[0].__int__(), momenta_before, momenta_after, ke_before, ke_after]
end make_row

table_data := flatten(make_row**(**elastic_total_momenta_before, **elastic_total_momenta_after, **elastic_total_ke_before, **elastic_total_ke_after))

elastic_calc_table := Table2D()

elastic_calc_table.set_labels(["Trial", "$p_i$", "$p_f$", "k_{ei}", "k_{ef}"])
elastic_calc_table.set_data(table_data)

"print(elastic_calc_table.latex())"


table_data := flatten(make_row**(**inelastic_total_momenta_before, **inelastic_total_momenta_after, **inelastic_total_ke_before, **inelastic_total_ke_after))
inelastic_calc_table := Table2D()

inelastic_calc_table.set_labels(["Trial", "$p_i$", "$p_f$", "k_{ei}", "k_{ef}"])
inelastic_calc_table.set_data(table_data)

print(inelastic_calc_table.latex())