import torch
import torch.nn as nn
import numpy as np

# --- 1. Setup and Constants ---
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Define constants (using symbolic representation for clarity)
hbar = 1.0  # Reduced Planck constant (set to 1 for simplicity)
T_total = 1.0 # Total simulation time

# --- 2. Define the Hamiltonian (The "Intelligent Problem") ---
class IntelligentHamiltonian(nn.Module):
    """
    This module defines the Hamiltonian H(x, t) for the system.
    For a "random but intelligent" problem, this is where the complexity lies.
    Example: A simple harmonic oscillator plus a time-dependent perturbation.
    """
    def __init__(self, num_states, initial_energy=1.0):
        super().__init__()
        self.num_states = num_states

        # State parameters (e.g., position/momentum representation)
        # These parameters define the system's intrinsic properties.
        self.position_params = nn.Parameter(torch.randn(num_states) * 0.1)
        self.potential_strength = nn.Parameter(torch.tensor(initial_energy))

    def forward(self, x):
        """
        Calculates the Hamiltonian H(x, t).
        x is the state vector (e.g., position/momentum).
        """
        # Kinetic Term (Simplified: depends on the state vector)
        # In a real system, this would involve derivatives of the state vector.
        kinetic_term = 0.5 * torch.sum(self.position_params**2, dim=0)

        # Potential Term (Time-dependent, making it "intelligent")
        # Example: A time-dependent potential that oscillates based on the state
        # The time dependence is crucial for the "intelligent" aspect.
        time_factor = torch.sin(torch.linspace(0, 2 * np.pi * T_total, self.num_states))
        potential_term = self.potential_strength * torch.cos(self.position_params) * time_factor

        H = kinetic_term + potential_term
        return H

# --- 3. The Schrödinger Solver (The Rheo Core) ---
class SchrödingerSolver(nn.Module):
    #Implements the time evolution using a time-splitting approach.
    #This mimics the continuous evolution $\frac{d\psi}{dt} = -i\hat{H}\psi$.
    
    def __init__(self, hamiltonian: IntelligentHamiltonian, num_states: int):
        super().__init__()
        self.H = hamiltonian.to(device)
        self.num_states = num_states

        # Initialize the wave function (Psi)
        # Psi shape: (Batch_size, Num_states)
        self.psi = torch.randn(1, num_states).to(device)

        # Define time steps (Crucial for stability, related to Stiffness Control)
        # A small dt is necessary for the Euler approximation to be valid.
        self.dt = T_total / 1000.0 # Small time step for integration
        self.num_steps = int(T_total / self.dt)
        print(f"Solver initialized with dt={self.dt:.6f} and {self.num_steps} steps.")

    def forward(self, t):
        """
        Performs the time evolution from t_start to t.
        This is the core integration loop.
        """
        # --- Rheo Step 1: Time Stepping ---
        current_psi = self.psi.clone()

        for step in range(self.num_steps):
            t_current = step * self.dt

            # 1. Calculate the Hamiltonian H at the current state
            H_t = self.H(current_psi)

            # 2. Apply the time evolution operator (i * dt * H)
            # psi(t+dt) = exp(-i * dt * H) * psi(t)
            # Using PyTorch's exp for the complex exponential: exp(A) = exp(Re(A)) * exp(Im(A))
            # The imaginary part is -dt * H
            evolution_operator = torch.exp(-1j * self.dt * H_t)

            current_psi = evolution_operator @ current_psi

            # --- Rheo Step 2: Error Constraint Check (Conceptual) ---
            # In a full Rheo system, we would calculate LTE here and adjust self.dt
            # if error_check(current_psi) > tolerance:
            #     self.dt *= 0.5 # Throttle down

        self.psi = current_psi
        return self.psi

# --- 4. Execution ---
if __name__ == "__main__":
    NUM_STATES = 10  # Dimension of the Hilbert space (e.g., number of basis functions)

    # 1. Initialize the Intelligent Hamiltonian
    H_model = IntelligentHamiltonian(num_states=NUM_STATES)
    H_model.eval() # Set to evaluation mode for fixed parameters during simulation

    # 2. Initialize the Schrödinger Solver
    solver = SchrödingerSolver(H_model, num_states=NUM_STATES)

    # 3. Run the simulation
    print("\nStarting time evolution...")
    final_psi = solver(t=T_total)

    print("\n--- Simulation Complete ---")
    print(f"Final Wave Function Shape: {final_psi.shape}")
    print("Sample of the final wave function (first state):")
    print(final_psi[0, :5])

    # Optional: Check if parameters have changed (they shouldn't if H_model.eval() is used)
    print("\nChecking Hamiltonian parameters (should be fixed):")
    print(f"Position Parameters Norm: {torch.norm(H_model.position_params).item():.4f}")