# -*- coding: utf-8 -*-
"""
Experimentul 1 — Delayed Choice Quantum Eraser (Kim, Yu, Kulik, Shih, Scully, PRL 84, 2000)
Baza analizei din superluminal.pdf (Lee Wen Wu, 2021), sectiunile 2-3.

Simulare cuantica locala (Monte Carlo pe amplitudini):
 - fotonul semnal cade pe ecranul D0 la pozitia x0, esantionata din |Psi_s(x)|^2;
 - starea fotonului idler colapseaza partial in alfa|A> + beta|B>,
   cu alfa = Psi_A(x0), beta = Psi_B(x0) (Eq. 11 din lucrare);
 - idlerul ajunge la D1/D2 (eraser, prin beam-splitter) sau D3/D4 (which-slit):
     P(D1) = |alfa+beta|^2/4   P(D2) = |alfa-beta|^2/4
     P(D3) = |alfa|^2/2        P(D4) = |beta|^2/2     (Eq. 22, 23, 14, 16)
 - sortarea pozitiilor x0 dupa detectorul idlerului reconstruieste R01..R04.

Concluzia fizica reprodusa: pe D0 TOTAL nu apar franje niciodata; franjele apar
numai la sortarea in coincidenta -> nu exista semnalizare (retro)cauzala.
"""
import sys
import numpy as np
import tkinter as tk
from tkinter import ttk
import matplotlib
matplotlib.use("TkAgg")
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

from phys_comun import amplitudini_fante, esantioneaza, contrast_franje, perioada_franje

XMAX = 12.0
NX = 2001
NBIN = 121


class AppDCQE:
    def __init__(self, root):
        self.root = root
        root.title("Exp. 1 — Delayed Choice Quantum Eraser (Kim et al. 2000)")
        self.rng = np.random.default_rng()
        self.x = np.linspace(-XMAX, XMAX, NX)
        self.bins = np.linspace(-XMAX, XMAX, NBIN + 1)
        self.centre = 0.5 * (self.bins[:-1] + self.bins[1:])
        self.ruleaza = False

        # ---- panou de control ----
        ctl = ttk.Frame(root, padding=8)
        ctl.pack(side=tk.LEFT, fill=tk.Y)
        ttk.Label(ctl, text="Delayed Choice\nQuantum Eraser",
                  font=("Segoe UI", 12, "bold")).pack(pady=(0, 8))
        self.btn = ttk.Button(ctl, text="Start", command=self.comuta)
        self.btn.pack(fill=tk.X, pady=2)
        ttk.Button(ctl, text="Reset", command=self.reset).pack(fill=tk.X, pady=2)

        ttk.Label(ctl, text="Fotoni / pas:").pack(anchor="w", pady=(10, 0))
        self.v_batch = tk.IntVar(value=400)
        ttk.Scale(ctl, from_=50, to=3000, variable=self.v_batch).pack(fill=tk.X)

        ttk.Label(ctl, text="Separare fante d [mm]:").pack(anchor="w", pady=(10, 0))
        self.v_d = tk.DoubleVar(value=0.35)
        ttk.Scale(ctl, from_=0.15, to=0.8, variable=self.v_d,
                  command=lambda e: self.reset()).pack(fill=tk.X)

        ttk.Label(ctl, text="Latime fanta a [mm]:").pack(anchor="w", pady=(10, 0))
        self.v_a = tk.DoubleVar(value=0.08)
        ttk.Scale(ctl, from_=0.03, to=0.2, variable=self.v_a,
                  command=lambda e: self.reset()).pack(fill=tk.X)

        self.lbl = ttk.Label(ctl, text="", justify="left", font=("Consolas", 9))
        self.lbl.pack(anchor="w", pady=12)

        ttk.Label(ctl, wraplength=200, justify="left", foreground="#555", text=(
            "R01/R02: idler la D1/D2 (eraser) -> franje complementare.\n"
            "R03/R04: idler la D3/D4 (which-slit) -> fara franje.\n"
            "D0 total = R01+R02+R03+R04 -> fara franje.\n\n"
            "Franjele exista doar in datele SORTATE in coincidenta, "
            "deci nu se poate transmite niciun semnal doar prin D0.")).pack(anchor="w")

        # ---- figura ----
        self.fig = Figure(figsize=(9, 7), dpi=100)
        gs = self.fig.add_gridspec(3, 2, hspace=0.5, wspace=0.3)
        self.ax0 = self.fig.add_subplot(gs[0, :])
        self.ax = [self.fig.add_subplot(gs[1, 0]), self.fig.add_subplot(gs[1, 1]),
                   self.fig.add_subplot(gs[2, 0]), self.fig.add_subplot(gs[2, 1])]
        self.canvas = FigureCanvasTkAgg(self.fig, master=root)
        self.canvas.get_tk_widget().pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)

        self.reset()
        if "--test" in sys.argv:
            self.comuta()
            root.after(2500, root.destroy)

    # ------------------------------------------------------------------
    def reset(self):
        self.psiA, self.psiB = amplitudini_fante(self.x, self.v_d.get(), self.v_a.get())
        self.p0 = (np.abs(self.psiA) ** 2 + np.abs(self.psiB) ** 2) / 2
        self.counts = {k: np.zeros(NBIN) for k in ("D0", "D1", "D2", "D3", "D4")}
        self.total = 0
        self.deseneaza()

    def comuta(self):
        self.ruleaza = not self.ruleaza
        self.btn.config(text="Pauza" if self.ruleaza else "Start")
        if self.ruleaza:
            self.pas()

    def pas(self):
        if not self.ruleaza:
            return
        n = int(self.v_batch.get())
        x0, idx = esantioneaza(self.x, self.p0, n, self.rng)
        al, be = self.psiA[idx], self.psiB[idx]
        s = np.abs(al) ** 2 + np.abs(be) ** 2
        # probabilitatile conditionate pentru detectorul idlerului
        p1 = np.abs(al + be) ** 2 / (4 * s)
        p2 = np.abs(al - be) ** 2 / (4 * s)
        p3 = np.abs(al) ** 2 / (2 * s)
        cum = np.cumsum(np.stack([p1, p2, p3], axis=1), axis=1)
        u = self.rng.random(n)[:, None]
        det = (u > cum).sum(axis=1)  # 0->D1, 1->D2, 2->D3, 3->D4
        self.counts["D0"] += np.histogram(x0, self.bins)[0]
        for j, nume in enumerate(("D1", "D2", "D3", "D4")):
            self.counts[nume] += np.histogram(x0[det == j], self.bins)[0]
        self.total += n
        self.deseneaza()
        self.root.after(30, self.pas)

    # ------------------------------------------------------------------
    def deseneaza(self):
        per = perioada_franje(self.v_d.get())
        self.ax0.clear()
        self.ax0.fill_between(self.centre, self.counts["D0"], color="#444", step="mid")
        self.ax0.set_title("D0 — toti fotonii semnal (ceea ce vede 'Bob' fara coincidente)")
        self.ax0.set_xlim(-XMAX, XMAX)

        culori = ("#1f77b4", "#d62728", "#2ca02c", "#9467bd")
        titluri = ("R01 (idler la D1 — eraser)", "R02 (idler la D2 — eraser)",
                   "R03 (idler la D3 — which-slit)", "R04 (idler la D4 — which-slit)")
        for axx, nume, cul, tit in zip(self.ax, ("D1", "D2", "D3", "D4"), culori, titluri):
            axx.clear()
            axx.fill_between(self.centre, self.counts[nume], color=cul, step="mid")
            axx.set_title(tit, fontsize=9)
            axx.set_xlim(-XMAX, XMAX)
        v01 = contrast_franje(self.counts["D1"], per, self.centre)
        v0 = contrast_franje(self.counts["D0"], per, self.centre)
        self.lbl.config(text=(f"fotoni:    {self.total}\n"
                              f"contrast R01: {v01:.3f}\n"
                              f"contrast D0:  {v0:.3f}"))
        self.canvas.draw_idle()


if __name__ == "__main__":
    root = tk.Tk()
    AppDCQE(root)
    root.mainloop()
