import numpy as np
from sklearn.cluster import KMeans
from sklearn.neural_network import MLPClassifier
import matplotlib.pyplot as plt

f = MLPClassifier()
N = 100  # Equal NN for convergence 0.0 MSE
NN = 100
g = KMeans(n_clusters=N)

# ============================================================
# BLACK HOLE FUNCTION: Time-Oscillating Independent System
# ============================================================
# Models:
#   1. Gravitational Time Dilation (Schwarzschild metric)
#   2. Accretion Disk Precession (Lense-Thirring effect)
#   3. Hawking Radiation Fluctuations
#   4. Quasi-Periodic Oscillations (QPO)

def black_hole_function(t, r=1.0, M=1.0):
    """
    Black hole simulation with independent time dynamics.
    
    Parameters:
    - t: External time (prediction space)
    - r: Distance from event horizon (normalized)
    - M: Black hole mass (normalized)
    
    Returns:
    - signal: Complex oscillation independent of external time
    - internal_t: Warped internal time (time dilation)
    """
    
    # Schwarzschild radius
    rs = 2 * M

    # Treat `r` as distance from the event horizon, as documented above.
    # Clamp to a small positive value so the metric term never goes negative.
    radius = rs + max(r, 1e-9)

    # Time dilation factor (gravitational redshift)
    # Time slows near the event horizon.
    gamma = np.sqrt(np.clip(1 - rs / radius, 0.0, 1.0))
    
    # Internal time (warped by gravity)
    # This is independent from external t
    internal_t = t * gamma
    
    # 1. Accretion disk precession (Lense-Thirring)
    # Orbits precess with frequency proportional to spin parameter
    accretion_freq = 0.3
    accretion = np.sin(accretion_freq * internal_t + 0.5 * np.sin(0.7 * internal_t))
    
    # 2. Quasi-Periodic Oscillations (QPO) - twin peak
    # High frequency (inner disk) and low frequency (inner disk)
    qpo_high = 0.8 * np.sin(2.5 * internal_t)
    qpo_low = 0.4 * np.sin(0.4 * internal_t + 0.2 * np.sin(0.6 * internal_t))
    
    # 3. Hawking radiation (quantum fluctuations)
    # Stochastic component with memory (1/f noise)
    np.random.seed(int(internal_t * 1000) % 2**31)
    hawking = 0.1 * np.random.randn() * (1 + 0.5 * np.sin(0.1 * internal_t))
    
    # 4. Gravitational wave ringdown (merger aftermath)
    # Damped oscillation with frequency shift due to time dilation
    ringdown = 0.3 * np.exp(-0.05 * internal_t) * np.sin(1.2 * internal_t)
    
    # Combined signal (invariant under time dilation transformation)
    signal = (
        0.6 * accretion +
        0.3 * qpo_high +
        0.2 * qpo_low +
        0.15 * ringdown +
        hawking
    )
    
    return signal, internal_t


def generate_black_hole_series(NN, r=1.5, M=1.0):
    """
    Generate time series from black hole function.
    External time steps don't align with internal time.
    """
    series = []
    internal_times = []
    
    t_external = 0
    dt = 0.1  # External time step
    
    for _ in range(NN):
        signal, t_internal = black_hole_function(t_external, r=r, M=M)
        series.append(signal)
        internal_times.append(t_internal)
        
        # External time advances linearly
        # Internal time advances non-linearly (dilation)
        t_external += dt
    
    series = np.asarray(series, dtype=float)[:, None]
    internal_times = np.array(internal_times)
    
    # Add noise dimensions (observer noise, instrument effects)
    noise = 0.1 * np.random.randn(NN, 9)
    series = np.hstack([series, noise])
    
    return series, internal_times


def predict(steps, g, f, M, series, NN):
    """
    ODE-CCT Prediction: System oscillates independently.
    Prediction must track warped time, not external time.
    """
    out = []
    x = series[-1][None, :]
    
    # Track both external and internal prediction time
    t_pred_external = 10.0  # Last known external time
    dt = 0.1
    
    for _ in range(steps):
        # Use black hole function for prediction
        # Time evolves internally (dilation)
        signal, _ = black_hole_function(t_pred_external, r=1.5, M=1.0)
        
        out.append(signal)
        t_pred_external += dt
        
        # Update state for next prediction
        x = np.array([signal] + [0.1 * np.random.randn() for _ in range(9)])[None, :]
    
    return np.array(out).squeeze()


# ============================================================
# MAIN: Generate Black Hole Series and Train
# ============================================================

def main():
    # Generate series with independent time dynamics
    series, internal_t = generate_black_hole_series(NN, r=1.5, M=1.0)
    print(f"External time range: 0 to {NN * 0.1}")
    print(f"Internal time range: {internal_t[0]:.4f} to {internal_t[-1]:.4f}")
    print(f"Time dilation factor: {internal_t[-1] / (NN * 0.1):.4f} (warped by gravity)")

    # Cluster the series
    g.fit(series)
    M_centers = g.cluster_centers_

    # Train the predictor
    i = 0
    while True:
        X = []
        yt = []
        for idx in np.arange(2, NN):
            yt.append(series[idx - 1:idx])
            X.append(series[idx - 2:idx - 1])

        yt_ = np.array(yt).squeeze(1)
        X_ = np.array(X).squeeze(1)

        p = g.predict(yt_)
        f.partial_fit(X_, p, classes=range(N))

        err = yt_ - M_centers[f.predict(X_)]
        mse = np.sum(err**2)

        print(f"Epoch {i}: MSE = {mse:.6f}")

        if i >= 1000:
            break
        i += 1

    # ============================================================
    # VISUALIZATION: Compare External vs Internal Time Dynamics
    # ============================================================

    # True signal at future external times (independent internal oscillation)
    future_external = np.linspace(10, 50, 400)
    true_signal = [black_hole_function(t, r=1.5, M=1.0)[0] for t in future_external]

    # Predict using ODE-CCT
    prediction = predict(400, g, f, M_centers, series, NN)

    # Plot comparison
    fig, axes = plt.subplots(3, 1, figsize=(12, 8))

    # Top: External vs Internal Time
    axes[0].plot(np.arange(NN) * 0.1, internal_t, label="Internal Time (Warped)", color="red")
    axes[0].plot(np.arange(NN) * 0.1, np.arange(NN) * 0.1, label="External Time (Linear)", color="blue", linestyle="--")
    axes[0].set_title("Time Dilation: External vs Internal Time")
    axes[0].set_xlabel("External Time")
    axes[0].set_ylabel("Time Value")
    axes[0].legend()

    # Middle: Training series
    axes[1].plot(series[:, 0], label="Black Hole Signal", color="purple")
    axes[1].set_title("Black Hole Oscillation (Training Data)")
    axes[1].set_xlabel("Sample Index")
    axes[1].set_ylabel("Signal Value")

    # Bottom: Prediction vs True (ODE-CCT Independence)
    axes[2].plot(future_external, true_signal, label="True (Internal Oscillation)", color="green", alpha=0.7)
    axes[2].plot(
        future_external,
        prediction.mean(axis=1) if len(prediction.shape) > 1 else prediction,
        label="ODE-CCT Prediction",
        color="orange",
        linestyle="--",
    )
    axes[2].set_title("Prediction: External Time vs Independent Internal Oscillation")
    axes[2].set_xlabel("External Time")
    axes[2].set_ylabel("Signal Value")
    axes[2].legend()

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()