"""
Modified DCQE experiment simulator based on:
Lee Wen Wu, "Superluminal Communication?" (March 2021)

This script reproduces the mathematical predictions described in sections 4.1-4.3
of the PDF, especially the author's claim that only the D5 path-combination setup
would generate a direct interference pattern on Bob's screen.
"""

from __future__ import annotations

import math
import queue
import threading
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

import numpy as np
import tkinter as tk
from tkinter import messagebox, scrolledtext

import matplotlib

matplotlib.use("TkAgg")
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.figure import Figure


BG = "#0b1020"
PANEL = "#121a2b"
CARD = "#182339"
ACCENT = "#2dd4bf"
ACCENT2 = "#38bdf8"
GREEN = "#22c55e"
AMBER = "#f59e0b"
RED = "#ef4444"
TEXT = "#e5eefb"
MUTED = "#94a3b8"
BORDER = "#24324a"


@dataclass
class ExperimentParameters:
    photons: int
    bins: int
    x_extent: float
    envelope_sigma: float
    fringe_frequency: float
    d5_phase_deg: float
    noise_fraction: float


@dataclass
class ScenarioResult:
    key: str
    title: str
    equation: str
    description: str
    direct_observation: bool
    theoretical_prob: np.ndarray
    sampled_counts: np.ndarray
    sampled_prob: np.ndarray
    baseline_contrast: float
    bob_sees_interference: bool
    peak_x: float


@dataclass
class SimulationReport:
    params: ExperimentParameters
    created_at: datetime
    x: np.ndarray
    baseline_prob: np.ndarray
    scenarios: list[ScenarioResult]
    summary: str


def normalize_prob(values: np.ndarray) -> np.ndarray:
    clipped = np.clip(np.asarray(values, dtype=float), 0.0, None)
    total = clipped.sum()
    if total <= 0:
        raise ValueError("Probability distribution collapsed to zero.")
    return clipped / total


def interference_contrast(prob: np.ndarray, baseline: np.ndarray) -> float:
    denom = float(np.max(baseline))
    if denom <= 0:
        return 0.0
    return float(np.max(np.abs(prob - baseline)) / denom)


def classify_visibility(contrast: float) -> tuple[bool, str]:
    if contrast >= 0.22:
        return True, "clear direct interference"
    if contrast >= 0.08:
        return True, "weak but visible interference"
    return False, "no direct interference"


def sample_distribution(prob: np.ndarray, photons: int, noise_fraction: float) -> tuple[np.ndarray, np.ndarray]:
    uniform = np.full_like(prob, 1.0 / len(prob))
    noisy = normalize_prob((1.0 - noise_fraction) * prob + noise_fraction * uniform)
    counts = np.random.multinomial(photons, noisy)
    return counts, counts / max(1, counts.sum())


def build_idler_amplitudes(params: ExperimentParameters) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    x = np.linspace(-params.x_extent, params.x_extent, params.bins)
    envelope = np.exp(-(x ** 2) / (2.0 * params.envelope_sigma ** 2))
    phase = params.fringe_frequency * x
    psi_a = envelope * np.exp(-0.5j * phase)
    psi_b = envelope * np.exp(+0.5j * phase)
    return x, psi_a, psi_b


def make_scenario(
    key: str,
    title: str,
    equation: str,
    description: str,
    direct_observation: bool,
    prob: np.ndarray,
    baseline: np.ndarray,
    x: np.ndarray,
    photons: int,
    noise_fraction: float,
) -> ScenarioResult:
    counts, sampled_prob = sample_distribution(prob, photons, noise_fraction)
    contrast = interference_contrast(prob, baseline)
    sees_interference, _ = classify_visibility(contrast)
    peak_x = float(x[int(np.argmax(prob))])
    return ScenarioResult(
        key=key,
        title=title,
        equation=equation,
        description=description,
        direct_observation=direct_observation,
        theoretical_prob=prob,
        sampled_counts=counts,
        sampled_prob=sampled_prob,
        baseline_contrast=contrast,
        bob_sees_interference=direct_observation and sees_interference,
        peak_x=peak_x,
    )


def simulate_pdf_experiment(params: ExperimentParameters, log_fn) -> SimulationReport:
    log_fn("Loading the PDF model from sections 4.1-4.3...")
    x, psi_a, psi_b = build_idler_amplitudes(params)

    p_a = normalize_prob(np.abs(psi_a) ** 2)
    p_b = normalize_prob(np.abs(psi_b) ** 2)
    p_mix = normalize_prob(0.5 * (p_a + p_b))

    log_fn("Building the four Bob-side distributions described in the paper...")

    # Eq. (30): Alice measures which-path with D3 / D4.
    p_which_path = p_mix

    # Eq. (32): Alice does nothing.
    p_no_action = p_mix

    # Eq. (37), (40), (41): quantum eraser with D1 / D2 and direct Bob screen.
    p_d1 = normalize_prob(0.5 * np.abs(psi_a + psi_b) ** 2)
    p_d2 = normalize_prob(0.5 * np.abs(psi_a - psi_b) ** 2)
    p_eraser_direct = normalize_prob(0.5 * p_d1 + 0.5 * p_d2)

    # Eq. (51): D5 path-combination hypothesis from the author.
    d5_phase = math.radians(params.d5_phase_deg)
    alpha = 1.0
    beta = np.exp(1j * d5_phase)
    p_d5 = normalize_prob(np.abs(alpha * psi_a + beta * psi_b) ** 2)

    scenarios = [
        make_scenario(
            key="which_path",
            title="Alice finds which-path information",
            equation="Eq. (30)",
            description="Alice uses D3 / D4. Bob sees the incoherent sum only.",
            direct_observation=True,
            prob=p_which_path,
            baseline=p_mix,
            x=x,
            photons=params.photons,
            noise_fraction=params.noise_fraction,
        ),
        make_scenario(
            key="no_action",
            title="Alice does nothing",
            equation="Eq. (32)",
            description="The pair remains entangled, but Bob still sees no direct fringes.",
            direct_observation=True,
            prob=p_no_action,
            baseline=p_mix,
            x=x,
            photons=params.photons,
            noise_fraction=params.noise_fraction,
        ),
        make_scenario(
            key="eraser_direct",
            title="Quantum eraser, Bob direct screen",
            equation="Eq. (41)",
            description="D1 / D2 fringes cancel in Bob's unsorted data.",
            direct_observation=True,
            prob=p_eraser_direct,
            baseline=p_mix,
            x=x,
            photons=params.photons,
            noise_fraction=params.noise_fraction,
        ),
        make_scenario(
            key="eraser_d1",
            title="Quantum eraser, coincidence on D1",
            equation="Eq. (37)",
            description="Interference appears only after coincidence sorting on D1.",
            direct_observation=False,
            prob=p_d1,
            baseline=p_mix,
            x=x,
            photons=max(1, params.photons // 2),
            noise_fraction=params.noise_fraction,
        ),
        make_scenario(
            key="eraser_d2",
            title="Quantum eraser, coincidence on D2",
            equation="Eq. (40)",
            description="The complementary fringe pattern with opposite phase.",
            direct_observation=False,
            prob=p_d2,
            baseline=p_mix,
            x=x,
            photons=max(1, params.photons // 2),
            noise_fraction=params.noise_fraction,
        ),
        make_scenario(
            key="d5",
            title="D5 path-combination hypothesis",
            equation="Eq. (51)",
            description="The author's proposed superluminal-signalling configuration.",
            direct_observation=True,
            prob=p_d5,
            baseline=p_mix,
            x=x,
            photons=params.photons,
            noise_fraction=params.noise_fraction,
        ),
    ]

    summary_lines = []
    for scenario in scenarios:
        if scenario.direct_observation:
            visible, label = classify_visibility(scenario.baseline_contrast)
            verdict = "YES" if scenario.bob_sees_interference else "NO"
            summary_lines.append(
                f"{scenario.title}: Bob direct interference = {verdict} "
                f"(contrast {scenario.baseline_contrast:.3f}, {label})"
            )
        else:
            summary_lines.append(
                f"{scenario.title}: conditioned pattern only "
                f"(contrast {scenario.baseline_contrast:.3f})"
            )

    log_fn("Simulation finished. Only the D5 mode should show direct fringes if the author's model is used.")
    return SimulationReport(
        params=params,
        created_at=datetime.now(),
        x=x,
        baseline_prob=p_mix,
        scenarios=scenarios,
        summary="\n".join(summary_lines),
    )


def export_results_txt(report: SimulationReport) -> Path:
    timestamp = report.created_at.strftime("%Y%m%d_%H%M%S")
    file_path = Path(__file__).resolve().with_name(f"superluminal_report_{timestamp}.txt")

    lines = [
        "Modified DCQE Report",
        "====================",
        "Source PDF: Lee Wen Wu, 'Superluminal Communication?' (March 2021)",
        f"Generated at: {report.created_at.isoformat(timespec='seconds')}",
        "",
        "Model note:",
        "This file reproduces the mathematical predictions from sections 4.1-4.3 of the PDF.",
        "It is a numerical model of the author's argument, not an experimental proof.",
        "",
        "Parameters",
        "----------",
        f"Photons per direct scenario: {report.params.photons}",
        f"Screen bins: {report.params.bins}",
        f"Screen half-width: {report.params.x_extent}",
        f"Envelope sigma: {report.params.envelope_sigma}",
        f"Fringe frequency: {report.params.fringe_frequency}",
        f"D5 phase (deg): {report.params.d5_phase_deg}",
        f"Noise fraction: {report.params.noise_fraction:.4f}",
        "",
        "Summary",
        "-------",
        report.summary,
        "",
        "Scenario details",
        "----------------",
    ]

    for scenario in report.scenarios:
        visibility_label = classify_visibility(scenario.baseline_contrast)[1]
        lines.extend(
            [
                f"Title: {scenario.title}",
                f"Paper equation: {scenario.equation}",
                f"Description: {scenario.description}",
                f"Bob direct observation: {'yes' if scenario.direct_observation else 'no, coincidence sorted only'}",
                f"Contrast vs incoherent baseline: {scenario.baseline_contrast:.6f}",
                f"Bob sees direct interference: {'yes' if scenario.bob_sees_interference else 'no'}",
                f"Visibility label: {visibility_label}",
                f"Peak screen x: {scenario.peak_x:.6f}",
                f"Sampled photons: {int(scenario.sampled_counts.sum())}",
                f"Sampled counts max bin: {int(scenario.sampled_counts.max())}",
                "",
            ]
        )

    file_path.write_text("\n".join(lines), encoding="utf-8")
    return file_path


def build_plots(fig: Figure, report: SimulationReport) -> None:
    fig.clear()
    gs = fig.add_gridspec(3, 2, hspace=0.38, wspace=0.22, left=0.06, right=0.97, top=0.93, bottom=0.07)
    order = ["which_path", "no_action", "eraser_direct", "eraser_d1", "eraser_d2", "d5"]
    scenario_map = {scenario.key: scenario for scenario in report.scenarios}
    x = report.x
    baseline = report.baseline_prob

    for idx, key in enumerate(order):
        scenario = scenario_map[key]
        ax = fig.add_subplot(gs[idx // 2, idx % 2])
        ax.plot(x, scenario.theoretical_prob, color=ACCENT2, linewidth=1.8, label="theory")
        ax.step(x, scenario.sampled_prob, where="mid", color=ACCENT, linewidth=1.0, alpha=0.9, label="sampled")
        ax.plot(x, baseline, color=MUTED, linestyle="--", linewidth=1.0, alpha=0.7, label="baseline")

        visible, label = classify_visibility(scenario.baseline_contrast)
        title_color = GREEN if scenario.bob_sees_interference else (ACCENT if visible else AMBER)
        ax.set_title(f"{scenario.title}\n{scenario.equation}", fontsize=8.5, color=title_color, pad=6)
        ax.text(
            0.02,
            0.94,
            f"contrast {scenario.baseline_contrast:.3f}\n{label}",
            transform=ax.transAxes,
            va="top",
            fontsize=7,
            color=TEXT,
            bbox={"facecolor": CARD, "edgecolor": BORDER, "boxstyle": "round,pad=0.3", "alpha": 0.9},
        )
        ax.set_xlim(x[0], x[-1])
        ax.set_ylim(0, max(0.001, scenario.theoretical_prob.max() * 1.18))
        ax.grid(alpha=0.22, color=BORDER)
        ax.tick_params(labelsize=7, colors=MUTED)
        if idx >= 4:
            ax.set_xlabel("Bob screen position x", fontsize=8, color=TEXT)
        if idx % 2 == 0:
            ax.set_ylabel("Probability", fontsize=8, color=TEXT)
        if idx == 0:
            ax.legend(fontsize=6.5, loc="upper right")

    fig.suptitle(
        "Modified DCQE reproduction from the PDF: only D5 should show direct fringes in the author's model",
        fontsize=10,
        color=TEXT,
    )
    fig.canvas.draw_idle()


class PDFExperimentApp(tk.Tk):
    def __init__(self) -> None:
        super().__init__()
        self.title("Modified DCQE Experiment")
        self.configure(bg=BG)
        self.geometry("1500x980")
        self.minsize(1180, 820)

        self._running = False
        self._ui_queue: queue.Queue[tuple] = queue.Queue()
        self._report: SimulationReport | None = None

        self._build_ui()
        self.after(50, self._process_ui_queue)

    def _build_ui(self) -> None:
        top = tk.Frame(self, bg=BG, pady=8)
        top.pack(fill="x", padx=12)

        tk.Label(top, text="MODIFIED DCQE EXPERIMENT", font=("Courier", 16, "bold"), bg=BG, fg=ACCENT).pack(
            side="left"
        )
        tk.Label(
            top,
            text="  Lee Wen Wu (2021)  |  sections 4.1-4.3",
            font=("Courier", 10),
            bg=BG,
            fg=MUTED,
        ).pack(side="left")

        cfg = tk.Frame(self, bg=PANEL, highlightthickness=1, highlightbackground=BORDER)
        cfg.pack(fill="x", padx=12, pady=(0, 8))
        inner = tk.Frame(cfg, bg=PANEL)
        inner.pack(fill="x", padx=10, pady=10)

        self.photons_var = tk.StringVar(value="40000")
        self.bins_var = tk.StringVar(value="260")
        self.x_extent_var = tk.StringVar(value="8.0")
        self.sigma_var = tk.StringVar(value="3.0")
        self.fringe_var = tk.StringVar(value="6.0")
        self.d5_phase_var = tk.StringVar(value="0.0")
        self.noise_var = tk.StringVar(value="1.5")

        self._make_entry(inner, 0, 0, "Photons / direct mode", self.photons_var, 10)
        self._make_entry(inner, 0, 2, "Screen bins", self.bins_var, 8)
        self._make_entry(inner, 0, 4, "Screen half-width", self.x_extent_var, 8)
        self._make_entry(inner, 1, 0, "Envelope sigma", self.sigma_var, 10)
        self._make_entry(inner, 1, 2, "Fringe frequency", self.fringe_var, 8)
        self._make_entry(inner, 1, 4, "D5 phase (deg)", self.d5_phase_var, 8)
        self._make_entry(inner, 1, 6, "Noise (%)", self.noise_var, 8)

        self.run_btn = tk.Button(
            inner,
            text="Run PDF model",
            font=("Courier", 10, "bold"),
            bg=ACCENT,
            fg=BG,
            relief="flat",
            cursor="hand2",
            padx=14,
            pady=5,
            command=self._start_experiment,
        )
        self.run_btn.grid(row=0, column=6, rowspan=1, padx=(18, 0), sticky="ew")

        info = tk.Label(
            inner,
            text="Direct Bob-screen comparisons: Eq.30, Eq.32, Eq.41, Eq.51",
            font=("Courier", 9),
            bg=PANEL,
            fg=MUTED,
        )
        info.grid(row=2, column=0, columnspan=7, sticky="w", pady=(10, 0))

        paned = tk.PanedWindow(self, orient="vertical", bg=BG, sashwidth=4, sashrelief="flat", sashpad=2)
        paned.pack(fill="both", expand=True, padx=12, pady=(0, 8))

        plot_frame = tk.Frame(paned, bg=PANEL, highlightthickness=1, highlightbackground=BORDER)
        paned.add(plot_frame, height=620, minsize=360)

        self.fig = Figure(figsize=(13, 7), dpi=96)
        self.canvas = FigureCanvasTkAgg(self.fig, master=plot_frame)
        self.canvas.get_tk_widget().pack(fill="both", expand=True)
        self._draw_placeholder()

        bottom = tk.Frame(paned, bg=BG)
        paned.add(bottom, minsize=180)

        self.cards_frame = tk.Frame(bottom, bg=BG)
        self.cards_frame.pack(fill="x", pady=(4, 6))

        log_frame = tk.Frame(bottom, bg=PANEL, highlightthickness=1, highlightbackground=BORDER)
        log_frame.pack(fill="both", expand=True)
        self.log = scrolledtext.ScrolledText(
            log_frame,
            height=9,
            font=("Courier", 8),
            bg=PANEL,
            fg=GREEN,
            insertbackground=GREEN,
            relief="flat",
            state="disabled",
            wrap="word",
        )
        self.log.pack(fill="both", expand=True, padx=4, pady=4)
        self.log.tag_config("error", foreground=RED)
        self.log.tag_config("accent", foreground=ACCENT)
        self.log.tag_config("amber", foreground=AMBER)
        self._log("Ready. This version reproduces the PDF model, not the Bell / IBM experiment.")

    def _make_entry(self, parent, row, column, label, variable, width) -> None:
        tk.Label(parent, text=label, font=("Courier", 9), bg=PANEL, fg=MUTED).grid(
            row=row,
            column=column,
            sticky="w",
            padx=(0, 6),
            pady=2,
        )
        entry = tk.Entry(
            parent,
            textvariable=variable,
            width=width,
            font=("Courier", 9),
            bg=CARD,
            fg=TEXT,
            relief="flat",
            highlightthickness=1,
            highlightbackground=BORDER,
            insertbackground=ACCENT,
        )
        entry.grid(row=row, column=column + 1, padx=(0, 14), pady=2, ipady=4)

    def _draw_placeholder(self) -> None:
        self.fig.clear()
        ax = self.fig.add_subplot(111)
        ax.set_facecolor(PANEL)
        ax.text(
            0.5,
            0.55,
            "Run the PDF model to compare Eq. (30), (32), (41) and (51).",
            transform=ax.transAxes,
            ha="center",
            va="center",
            fontsize=13,
            color=MUTED,
        )
        ax.text(
            0.5,
            0.45,
            "The author claims that only the D5 configuration should show direct fringes on Bob's screen.",
            transform=ax.transAxes,
            ha="center",
            va="center",
            fontsize=10,
            color=MUTED,
        )
        ax.axis("off")
        self.canvas.draw()

    def _log(self, msg: str, tag: str | None = None) -> None:
        if threading.current_thread() is not threading.main_thread():
            self._ui_queue.put(("log", msg, tag))
            return

        self.log.config(state="normal")
        line = f"[{datetime.now().strftime('%H:%M:%S')}] {msg}\n"
        if tag:
            self.log.insert("end", line, tag)
        else:
            self.log.insert("end", line)
        self.log.see("end")
        self.log.config(state="disabled")
        self.update_idletasks()

    def _process_ui_queue(self) -> None:
        try:
            while True:
                action, *payload = self._ui_queue.get_nowait()
                if action == "log":
                    self._log(payload[0], payload[1])
                elif action == "report":
                    report = payload[0]
                    self._apply_report(report)
                elif action == "finished":
                    self._running = False
                    self.run_btn.config(state="normal", text="Run PDF model")
        except queue.Empty:
            pass
        except tk.TclError:
            return

        try:
            self.after(50, self._process_ui_queue)
        except tk.TclError:
            pass

    def _parse_params(self) -> ExperimentParameters:
        try:
            photons = int(self.photons_var.get())
            bins = int(self.bins_var.get())
            x_extent = float(self.x_extent_var.get())
            sigma = float(self.sigma_var.get())
            fringe = float(self.fringe_var.get())
            d5_phase = float(self.d5_phase_var.get())
            noise_percent = float(self.noise_var.get())
        except ValueError as exc:
            raise ValueError("All parameters must be numeric.") from exc

        if photons < 100:
            raise ValueError("Photons must be at least 100.")
        if bins < 80:
            raise ValueError("Screen bins must be at least 80.")
        if x_extent <= 0 or sigma <= 0:
            raise ValueError("Screen size and envelope sigma must be positive.")
        if fringe <= 0:
            raise ValueError("Fringe frequency must be positive.")
        if not 0.0 <= noise_percent <= 50.0:
            raise ValueError("Noise must stay between 0 and 50 percent.")

        return ExperimentParameters(
            photons=photons,
            bins=bins,
            x_extent=x_extent,
            envelope_sigma=sigma,
            fringe_frequency=fringe,
            d5_phase_deg=d5_phase,
            noise_fraction=noise_percent / 100.0,
        )

    def _start_experiment(self) -> None:
        if self._running:
            return

        try:
            params = self._parse_params()
        except ValueError as exc:
            messagebox.showerror("Invalid parameters", str(exc))
            return

        self._running = True
        self.run_btn.config(state="disabled", text="Running...")
        threading.Thread(target=self._run_experiment, args=(params,), daemon=True).start()

    def _run_experiment(self, params: ExperimentParameters) -> None:
        try:
            self._log("=" * 60, "accent")
            self._log("Starting modified DCQE reproduction from the PDF...", "accent")
            self._log(
                f"photons={params.photons}, bins={params.bins}, sigma={params.envelope_sigma}, "
                f"fringe={params.fringe_frequency}, D5 phase={params.d5_phase_deg} deg"
            )
            report = simulate_pdf_experiment(params, self._log)
            report_path = export_results_txt(report)
            self._log(f"TXT report saved: {report_path}", "accent")
            self._log(report.summary, "amber")
            self._log("=" * 60, "accent")
            self._ui_queue.put(("report", report))
        except Exception as exc:  # noqa: BLE001
            self._log(f"Error: {exc}", "error")
        finally:
            self._ui_queue.put(("finished",))

    def _apply_report(self, report: SimulationReport) -> None:
        self._report = report
        build_plots(self.fig, report)
        self._rebuild_cards(report)
        self.canvas.draw()

    def _rebuild_cards(self, report: SimulationReport) -> None:
        for widget in self.cards_frame.winfo_children():
            widget.destroy()

        direct_keys = ["which_path", "no_action", "eraser_direct", "d5", "eraser_d1", "eraser_d2"]
        scenario_map = {scenario.key: scenario for scenario in report.scenarios}

        for idx, key in enumerate(direct_keys):
            scenario = scenario_map[key]
            visible, label = classify_visibility(scenario.baseline_contrast)
            if scenario.bob_sees_interference:
                border = GREEN
                verdict = "Bob direct: interference"
            elif scenario.direct_observation:
                border = AMBER
                verdict = "Bob direct: no fringes"
            else:
                border = ACCENT2
                verdict = "Coincidence sorted only"

            card = tk.Frame(
                self.cards_frame,
                bg=CARD,
                padx=10,
                pady=8,
                highlightthickness=1,
                highlightbackground=border,
            )
            card.grid(row=idx // 3, column=idx % 3, padx=6, pady=4, sticky="nsew")

            tk.Label(card, text=scenario.title, font=("Courier", 8, "bold"), bg=CARD, fg=border, wraplength=320).pack(
                anchor="w"
            )
            tk.Label(card, text=scenario.equation, font=("Courier", 8), bg=CARD, fg=MUTED).pack(anchor="w")
            tk.Label(card, text=verdict, font=("Courier", 8), bg=CARD, fg=TEXT).pack(anchor="w")
            tk.Label(
                card,
                text=f"contrast={scenario.baseline_contrast:.3f} | {label}",
                font=("Courier", 7),
                bg=CARD,
                fg=MUTED,
            ).pack(anchor="w")
            tk.Label(card, text=scenario.description, font=("Courier", 7), bg=CARD, fg=MUTED, wraplength=320).pack(
                anchor="w"
            )

        for column in range(3):
            self.cards_frame.columnconfigure(column, weight=1)


if __name__ == "__main__":
    app = PDFExperimentApp()
    app.mainloop()