
import numpy as np
from scipy.linalg import expm
from scipy.optimize import minimize

def trotter_evolution(hamiltonian: np.ndarray, initial_state: np.ndarray,
                      time: float, n_steps: int = 10) -> dict:
    """
    First-order Trotterization: exp(-iHt) ~ [exp(-iHt/n)]^n
    Works for any Hamiltonian passed as a dense matrix.
    """
    dt     = time / n_steps
    U_step = expm(-1j * hamiltonian * dt)
    state  = initial_state.copy().astype(complex)
    states = [state.copy()]
    for _ in range(n_steps):
        state = U_step @ state
        states.append(state.copy())
    exact_state = expm(-1j * hamiltonian * time) @ initial_state
    fidelity    = float(np.abs(np.conj(exact_state) @ state) ** 2)
    return {
        "final_state":  state,
        "exact_state":  exact_state,
        "fidelity":     fidelity,
        "trajectory":   states,
        "time":         time,
        "n_steps":      n_steps,
        "trotter_error": float(np.linalg.norm(state - exact_state)),
    }

def varqrte(hamiltonian: np.ndarray, initial_state: np.ndarray,
            time: float, n_steps: int = 10, n_params: int = None) -> dict:
    """
    Variational Real-Time Evolution (VarQRTE).
    Fits ansatz parameters to follow real-time evolution.
    """
    dim      = len(initial_state)
    n_params = n_params or dim
    params   = np.zeros(n_params)

    def ansatz(p):
        angles  = p[:dim] if len(p) >= dim else np.pad(p, (0, dim - len(p)))
        phases  = np.exp(1j * angles)
        s       = initial_state * phases
        s      /= np.linalg.norm(s)
        return s

    exact_final = expm(-1j * hamiltonian * time) @ initial_state
    def loss(p):
        s = ansatz(p)
        return 1 - float(np.abs(np.conj(exact_final) @ s) ** 2)

    res    = minimize(loss, params, method="COBYLA", options={"maxiter": 300})
    approx = ansatz(res.x)
    return {
        "final_state":    approx,
        "exact_state":    exact_final,
        "fidelity":       1 - res.fun,
        "optimal_params": res.x,
        "converged":      res.success,
        "method":         "VarQRTE",
    }

def varqite(hamiltonian: np.ndarray, initial_state: np.ndarray,
            time: float, n_steps: int = 10) -> dict:
    """
    Variational Imaginary-Time Evolution (VarQITE).
    Evolves under exp(-H*tau) (imaginary time) to find ground state.
    """
    exact_imag = expm(-hamiltonian * time) @ initial_state
    norm       = np.linalg.norm(exact_imag)
    exact_imag = exact_imag / norm if norm > 0 else exact_imag

    dim    = len(initial_state)
    params = np.zeros(dim)
    def ansatz(p):
        phases = np.exp(1j * p[:dim] if len(p) >= dim else np.pad(p, (0, dim - len(p))))
        s      = initial_state * phases
        return s / np.linalg.norm(s)
    def loss(p):
        return 1 - float(np.abs(np.conj(exact_imag) @ ansatz(p)) ** 2)

    res = minimize(loss, params, method="COBYLA", options={"maxiter": 300})
    evals = np.linalg.eigvalsh(hamiltonian)
    return {
        "final_state":    ansatz(res.x),
        "exact_state":    exact_imag,
        "fidelity":       1 - res.fun,
        "ground_energy":  float(evals[0]),
        "converged":      res.success,
        "method":         "VarQITE",
    }
