from fractions import Fraction as Q
from math import cos, sin, pi
from komorebi_browser import show_plot

# A = [[0, -1], [1, 0]], psi(0) = (1, 1).
# psi is intentionally unnormalized. The conserved squared norm is 2.
x, y = Q(1), Q(1)
incoming = (2*x*(-y), 2*y*x)
assert incoming == (-2, 2) and sum(incoming) == 0
print(f"Initial packet squared norms: {x*x}, {y*y}")
print(f"Initial incoming currents: {incoming[0]}, {incoming[1]}")
print("Total squared norm:", x*x+y*y)

times = [2*pi*j/160 for j in range(161)]
e0 = [(cos(t)-sin(t))**2 for t in times]
e1 = [(sin(t)+cos(t))**2 for t in times]
show_plot([
    {"x": times, "y": e0, "name": "Packet 0", "mode": "lines"},
    {"x": times, "y": e1, "name": "Packet 1", "mode": "lines", "line": {"dash": "dash"}},
    {"x": times, "y": [a+b for a,b in zip(e0,e1)], "name": "Total", "mode": "lines", "line": {"dash": "dot"}},
], {"title": {"text": "Exchange between two orthogonal packets"},
    "xaxis": {"title": {"text": "Time t"}},
    "yaxis": {"title": {"text": "Squared norm"}}},
    description="The two packet squared norms exchange between zero and two. Their sum stays at two. Divide each by two to get normalized probabilities. The initial incoming currents are minus two and plus two.")
print("Plotted 161 times on [0, 2 pi].")