# SPDX-License-Identifier: GPL-3.0-or-later
"""
Temporal Phase-Offset Brane Cosmology

Copyright (C) 2026 Christian Sacks

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""

import matplotlib
matplotlib.use("TkAgg")
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Slider

# ------------------- Galaxy configuration -------------------
num_stars = 2000
num_arms = 2
max_radius = 5.0

stars_per_arm = num_stars // num_arms
r = np.sqrt(np.random.rand(num_stars)) * max_radius
arm_indices = np.repeat(np.arange(num_arms), stars_per_arm)

# Spiral arms (k=-2.0 for trailing arms)
k = -2.0  
initial_theta = k * r + np.random.normal(0, 0.2, num_stars)
initial_theta += (2*np.pi / num_arms) * arm_indices
current_angles = initial_theta.copy()

# ------------------- Brane configuration -------------------
x_brane = np.linspace(-10, 10, 400)
time_val = 0.0
phase_speed = 0.01 
temporal_offset_enabled = True
amp = 6.0 

# ------------------- Figure setup -------------------
fig, (ax_brane, ax_galaxy) = plt.subplots(1, 2, figsize=(14, 6))
plt.subplots_adjust(bottom=0.25)

# 3-Brane Plot (Bulk Physics)
ax_brane.set_xlim(-10, 10)
ax_brane.set_ylim(-8, 8) 
ax_brane.set_title("5D Bulk: Temporal Shear S_ij")
ax_brane.set_ylabel("Warp Dimension (ξ)")
ax_brane.set_facecolor('#1a1a1a')

brane_top, = ax_brane.plot(x_brane, np.ones_like(x_brane)*4, lw=1.5, color='cyan', alpha=0.5, label="Confining 1")
brane_mid, = ax_brane.plot(x_brane, np.zeros_like(x_brane), lw=2.5, color='white', label="Our Brane (0)")
brane_bot, = ax_brane.plot(x_brane, np.ones_like(x_brane)*-4, lw=1.5, color='cyan', alpha=0.5, label="Confining 2")

# Shear visualization (vertical lines between branes)
shear_fill = ax_brane.fill_between(x_brane, 0, 0, color='magenta', alpha=0.2, label="Shear Field")
ax_brane.legend(loc="upper right", fontsize='small')

# Text meter for real-time values
shear_text = ax_brane.text(-9.5, 7, '', color='magenta', fontweight='bold')

# Galaxy Plot
ax_galaxy.set_xlim(-6, 6)
ax_galaxy.set_ylim(-6, 6)
ax_galaxy.set_aspect('equal')
ax_galaxy.set_facecolor('black')
ax_galaxy.set_title("Observable Universe (Galaxy)")
core = plt.Circle((0, 0), 0.3, color='yellow', zorder=10)
ax_galaxy.add_artist(core)

ref_line, = ax_galaxy.plot([0, max_radius], [0, 0], color='red', lw=3, alpha=0.6, linestyle='--', zorder=5)
star_scat = ax_galaxy.scatter([], [], color='white', s=1, alpha=0.8)

# ------------------- Slider & Keys -------------------
ax_slider = plt.axes([0.2, 0.1, 0.65, 0.03], facecolor='gray')
slider = Slider(ax_slider, 'Bulk Tension (λ)', 0.0, 6.0, valinit=amp)

paused = False 

def on_key(event):
    global temporal_offset_enabled, paused, current_angles, time_val
    if event.key == 't':
        temporal_offset_enabled = not temporal_offset_enabled
    elif event.key == 'p':
        paused = not paused
        if paused: ani.event_source.stop()
        else: ani.event_source.start()
    elif event.key == 'r':
        time_val = 0.0
        current_angles = initial_theta.copy()
    elif event.key == '.':
        slider.set_val(slider.valmax)
        temporal_offset_enabled = True

    # Arrow key slider control
    step = 1.0 if 'shift' in event.key else 0.1
    if 'right' in event.key:
        slider.set_val(min(slider.val + step, slider.valmax))
    elif 'left' in event.key:
        slider.set_val(max(slider.val - step, slider.valmin))

fig.canvas.mpl_connect('key_press_event', on_key)

# ------------------- Animation -------------------
def update(frame):
    global time_val, current_angles, shear_fill
    time_val += phase_speed
    amp_val = slider.val

    # 1. Update Branes & Calculate Shear
    # t_i = t + f_i(x_perp) [cite: 19]
    wave = amp_val * np.sin(0.8 * x_brane + time_val)
    
    if temporal_offset_enabled:
        y_top = 4 + wave * 0.5
        y_bot = -4 - wave * 0.5
        y_mid = wave * 0.2
        
        # Calculate instantaneous shear S_ij as the displacement 
        # In this model, max shear is the amplitude of the offset
        current_shear = np.abs(y_top[0] - y_mid[0])
        shear_text.set_text(f'Temporal Shear S_ij: {current_shear:.2f}')
    else:
        y_top, y_bot, y_mid = np.ones_like(x_brane)*4, np.ones_like(x_brane)*-4, np.zeros_like(x_brane)
        shear_text.set_text('Shear: 0.00')

    brane_top.set_ydata(y_top)
    brane_bot.set_ydata(y_bot)
    brane_mid.set_ydata(y_mid)
    
    # Update fill to visualize the 'Radion Field' / distance between branes 
    shear_fill.remove()
    shear_fill = ax_brane.fill_between(x_brane, y_mid, y_top, color='magenta', alpha=0.1)

    # 2. Galaxy Physics
    omega_anchor = 0.04 
    omega_kepler = omega_anchor / (1 + r*0.6)

    if temporal_offset_enabled:
        influence = (amp_val / slider.valmax) 
        omega = (1 - influence) * omega_kepler + (influence * omega_anchor)
    else:
        omega = omega_kepler

    current_angles += omega
    star_scat.set_offsets(np.c_[r * np.cos(current_angles), r * np.sin(current_angles)])

    ref_angle = (frame * omega_anchor) + (np.pi/2)
    ref_line.set_data([0, max_radius * np.cos(ref_angle)], [0, max_radius * np.sin(ref_angle)])

    return brane_top, brane_mid, brane_bot, star_scat, ref_line, shear_text

ani = FuncAnimation(fig, update, interval=30, cache_frame_data=False, blit=False)
plt.show()