# -*- coding: utf-8 -*-
"""
Extensión del simulador Lindblad para manejar Hamiltonianos dependientes del tiempo.
Combina evolución coherente (Schrödinger) con procesos disipativos (Lindblad).

Estrategias implementadas:
1. Hamiltoniano parametrizado H(t, params)
2. Operadores de Lindblad también dependientes del tiempo L_k(t)
3. Integración adaptativa con dependencia temporal explícita
4. Ejemplos: pulsos láser, campos oscilantes, rampas adiabáticas

Autor: Extension del código base de Jacobo Tlacaelel Mina Rodriguez
"""

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from typing import Callable, List, Optional, Dict, Any
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)

@dataclass
class TimeDependentSystemParameters:
    """Parámetros para sistemas con dependencia temporal."""
    name: str
    dimension: int
    H_func: Callable[[float, Dict], np.ndarray]  # H(t, params)
    L_func: Callable[[float, Dict], List[np.ndarray]]  # L_operators(t, params)
    rho_0: np.ndarray
    t_span: tuple
    time_params: Dict[str, Any]  # Parámetros para funciones temporales
    observables: Dict[str, np.ndarray]

class TimeDependentLindladSimulator:
    """Simulador Lindblad con Hamiltoniano dependiente del tiempo."""
    
    def __init__(self, config=None):
        self.config = config
        self.results = {}
    
    def lindblad_rhs_time_dependent(self, t: float, rho_vec: np.ndarray,
                                  H_func: Callable, L_func: Callable, 
                                  params: Dict) -> np.ndarray:
        """
        Ecuación de Lindblad con dependencia temporal explícita.
        
        drho/dt = -i[H(t), rho] + Σ_k (L_k(t) rho L_k†(t) - 1/2 {L_k†(t)L_k(t), rho})
        """
        try:
            n = int(np.sqrt(len(rho_vec)))
            rho = rho_vec.reshape((n, n))
            
            # Hamiltoniano en el tiempo t
            H_t = H_func(t, params)
            
            # Operadores de Lindblad en el tiempo t
            L_operators_t = L_func(t, params)
            
            # Término coherente: -i[H(t), rho]
            drho_dt = -1j * (H_t @ rho - rho @ H_t)
            
            # Términos disipativos
            for L_t in L_operators_t:
                L_dagger_t = L_t.conj().T
                drho_dt += (L_t @ rho @ L_dagger_t - 
                           0.5 * (L_dagger_t @ L_t @ rho + rho @ L_dagger_t @ L_t))
            
            return drho_dt.flatten()
            
        except Exception as e:
            logger.error(f"Error en lindblad_rhs_time_dependent en t={t}: {e}")
            raise
    
    def simulate(self, params: TimeDependentSystemParameters, 
                t_eval: Optional[np.ndarray] = None) -> Dict:
        """Simula sistema con dependencia temporal."""
        
        if t_eval is None:
            t_eval = np.linspace(params.t_span[0], params.t_span[1], 200)
        
        try:
            sol = solve_ivp(
                fun=self.lindblad_rhs_time_dependent,
                t_span=params.t_span,
                y0=params.rho_0.flatten(),
                args=(params.H_func, params.L_func, params.time_params),
                t_eval=t_eval,
                rtol=1e-8,
                atol=1e-10,
                method='RK45'
            )
            
            if not sol.success:
                raise RuntimeError(f"Integración falló: {sol.message}")
            
            # Procesar resultados
            rho_t = np.array([s.reshape((params.dimension, params.dimension)) 
                            for s in sol.y.T])
            
            # Calcular observables
            observables = {}
            for obs_name, obs_op in params.observables.items():
                observables[obs_name] = [np.trace(rho @ obs_op) for rho in rho_t]
            
            # Métricas
            traces = [np.trace(rho).real for rho in rho_t]
            purities = [np.trace(rho @ rho).real for rho in rho_t]
            
            results = {
                'success': True,
                'name': params.name,
                'time': sol.t,
                'rho_t': rho_t,
                'observables': observables,
                'traces': traces,
                'purities': purities,
                'H_evolution': [params.H_func(t, params.time_params) for t in sol.t],
                'L_evolution': [params.L_func(t, params.time_params) for t in sol.t]
            }
            
            return results
            
        except Exception as e:
            logger.error(f"Error en simulación: {e}")
            return {'success': False, 'error': str(e)}

# ============================================================================
# EJEMPLOS DE SISTEMAS CON DEPENDENCIA TEMPORAL
# ============================================================================

def create_driven_qubit_system():
    """
    Qubit controlado por pulsos láser + decoherencia.
    
    H(t) = ω₀/2 σz + Ω(t)/2 σx    # Campo de control
    L(t) = √γ(t) σ₋               # Decaimiento variable
    """
    
    # Operadores base
    sigma_z = np.array([[1, 0], [0, -1]], dtype=complex)
    sigma_x = np.array([[0, 1], [1, 0]], dtype=complex)
    sigma_minus = np.array([[0, 0], [1, 0]], dtype=complex)
    
    def H_driven_qubit(t, params):
        """Hamiltoniano con pulso Gaussiano."""
        omega_0 = params['omega_0']
        omega_drive = params['omega_drive']
        pulse_duration = params['pulse_duration']
        pulse_center = params['pulse_center']
        
        # Pulso Gaussiano
        pulse_envelope = np.exp(-((t - pulse_center)/pulse_duration)**2)
        Omega_t = omega_drive * pulse_envelope
        
        return 0.5 * omega_0 * sigma_z + 0.5 * Omega_t * sigma_x
    
    def L_driven_qubit(t, params):
        """Operadores de Lindblad con decay modulado."""
        gamma_base = params['gamma_base']
        
        # Ejemplo: decaimiento que aumenta durante el pulso
        pulse_center = params['pulse_center']
        pulse_duration = params['pulse_duration']
        pulse_factor = 1 + 0.5 * np.exp(-((t - pulse_center)/pulse_duration)**2)
        
        gamma_t = gamma_base * pulse_factor
        return [np.sqrt(gamma_t) * sigma_minus]
    
    # Parámetros temporales
    time_params = {
        'omega_0': 1.0,
        'omega_drive': 2.0,
        'pulse_duration': 1.0,
        'pulse_center': 5.0,
        'gamma_base': 0.1
    }
    
    # Estado inicial (ground state)
    rho_0 = np.array([[1, 0], [0, 0]], dtype=complex)
    
    # Observables
    observables = {
        'sigma_z': sigma_z,
        'sigma_x': sigma_x,
        'population_excited': np.array([[0, 0], [0, 1]], dtype=complex)
    }
    
    return TimeDependentSystemParameters(
        name="driven_qubit",
        dimension=2,
        H_func=H_driven_qubit,
        L_func=L_driven_qubit,
        rho_0=rho_0,
        t_span=(0, 10),
        time_params=time_params,
        observables=observables
    )

def create_adiabatic_passage_system():
    """
    Transferencia adiabática poblacional estimulada (STIRAP).
    Sistema de 3 niveles con dos campos láser contrapropagantes.
    
    |1⟩ ←→ |2⟩ ←→ |3⟩
       Ω₁(t)   Ω₂(t)
    """
    
    # Operadores de transición
    sigma_12 = np.array([[0, 1, 0], [0, 0, 0], [0, 0, 0]], dtype=complex)
    sigma_23 = np.array([[0, 0, 0], [0, 0, 1], [0, 0, 0]], dtype=complex)
    sigma_13 = np.array([[0, 0, 1], [0, 0, 0], [0, 0, 0]], dtype=complex)
    
    # Proyectores
    P1 = np.array([[1, 0, 0], [0, 0, 0], [0, 0, 0]], dtype=complex)
    P2 = np.array([[0, 0, 0], [0, 1, 0], [0, 0, 0]], dtype=complex)
    P3 = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 1]], dtype=complex)
    
    def H_stirap(t, params):
        """Hamiltoniano STIRAP con pulsos contrapropagantes."""
        delta = params['detuning']
        Omega_max = params['Omega_max']
        T_pulse = params['T_pulse']
        t_delay = params['t_delay']
        
        # Pulsos Gaussianos con retraso
        Omega1_t = Omega_max * np.exp(-((t - T_pulse)/T_pulse * 2)**2)
        Omega2_t = Omega_max * np.exp(-((t - T_pulse - t_delay)/T_pulse * 2)**2)
        
        H = delta * P2  # Desintonía del nivel intermedio
        H += 0.5 * Omega1_t * (sigma_12 + sigma_12.conj().T)  # Acoplamiento 1↔2
        H += 0.5 * Omega2_t * (sigma_23 + sigma_23.conj().T)  # Acoplamiento 2↔3
        
        return H
    
    def L_stirap(t, params):
        """Decaimiento espontáneo desde el nivel excitado."""
        gamma2 = params['gamma2']
        
        # Solo el nivel 2 decae
        L1 = np.sqrt(gamma2) * (sigma_12.conj().T)  # 2 → 1
        L3 = np.sqrt(gamma2) * (sigma_23.conj().T)  # 2 → 3
        
        return [L1, L3]
    
    time_params = {
        'detuning': 0.0,      # Resonancia
        'Omega_max': 1.0,     # Frecuencia de Rabi máxima
        'T_pulse': 2.0,       # Duración del pulso
        't_delay': 1.0,       # Retraso entre pulsos (crucial para STIRAP)
        'gamma2': 0.05        # Decaimiento del nivel intermedio
    }
    
    # Estado inicial: |1⟩
    rho_0 = P1.copy()
    
    observables = {
        'Population_1': P1,
        'Population_2': P2, 
        'Population_3': P3,
        'Coherence_13': sigma_13 + sigma_13.conj().T
    }
    
    return TimeDependentSystemParameters(
        name="STIRAP_3level",
        dimension=3,
        H_func=H_stirap,
        L_func=L_stirap,
        rho_0=rho_0,
        t_span=(0, 8),
        time_params=time_params,
        observables=observables
    )

def create_parametric_oscillator():
    """
    Oscilador paramétrico con bomba modulada + decoherencia.
    
    H(t) = ω a†a + f(t)(a² + a†²)    # Bomba paramétrica
    L = √κ a + √γφ (a + a†)          # Pérdidas + dephasing de fase
    """
    N = 6  # Dimensión del espacio de Fock
    
    # Operadores
    a = np.diag(np.sqrt(np.arange(1, N)), k=1).astype(complex)
    adagger = a.conj().T
    n_op = adagger @ a
    
    # Operadores cuadráticos
    a_squared = a @ a
    adagger_squared = adagger @ adagger
    
    def H_parametric(t, params):
        """Hamiltoniano con bomba paramétrica modulada."""
        omega = params['omega']
        f_max = params['f_max']
        f_freq = params['f_freq']
        
        # Modulación sinusoidal de la bomba
        f_t = f_max * np.sin(f_freq * t)
        
        H = omega * n_op + f_t * (a_squared + adagger_squared)
        return H
    
    def L_parametric(t, params):
        """Pérdidas de cavidad + dephasing de fase."""
        kappa = params['kappa']
        gamma_phi = params['gamma_phi']
        
        L_loss = np.sqrt(kappa) * a
        L_dephase = np.sqrt(gamma_phi) * (a + adagger)  # Dephasing de cuadratura
        
        return [L_loss, L_dephase]
    
    time_params = {
        'omega': 1.0,
        'f_max': 0.5,         # Amplitud de bomba paramétrica
        'f_freq': 2.0,        # Frecuencia de modulación
        'kappa': 0.1,         # Pérdidas de cavidad
        'gamma_phi': 0.05     # Dephasing de fase
    }
    
    # Estado inicial: vacío
    rho_0 = np.zeros((N, N), dtype=complex)
    rho_0[0, 0] = 1.0
    
    # Observables
    x_op = (a + adagger) / np.sqrt(2)      # Posición
    p_op = -1j * (a - adagger) / np.sqrt(2) # Momento
    
    observables = {
        'photon_number': n_op,
        'position': x_op,
        'momentum': p_op,
        'squeezing_x': x_op @ x_op,
        'squeezing_p': p_op @ p_op
    }
    
    return TimeDependentSystemParameters(
        name="parametric_oscillator",
        dimension=N,
        H_func=H_parametric,
        L_func=L_parametric,
        rho_0=rho_0,
        t_span=(0, 20),
        time_params=time_params,
        observables=observables
    )

def plot_time_dependent_results(results: Dict, save_path: Optional[str] = None):
    """Visualización especializada para sistemas con dependencia temporal."""
    
    if not results.get('success', False):
        print("No se pueden graficar resultados fallidos")
        return
    
    time = results['time']
    observables = results['observables']
    
    fig = plt.figure(figsize=(15, 10))
    
    # Layout de subplots
    gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)
    
    # 1. Evolución de observables principales
    ax1 = fig.add_subplot(gs[0, :2])
    for obs_name, obs_vals in observables.items():
        if 'Population' in obs_name or 'population' in obs_name:
            ax1.plot(time, np.real(obs_vals), 'o-', label=obs_name, markersize=3)
    ax1.set_xlabel('Tiempo')
    ax1.set_ylabel('Población')
    ax1.set_title('Evolución de Poblaciones')
    ax1.legend()
    ax1.grid(True)
    
    # 2. Conservación de traza y pureza
    ax2 = fig.add_subplot(gs[0, 2])
    ax2.plot(time, results['traces'], 'b-', label='Traza')
    ax2.plot(time, results['purities'], 'r--', label='Pureza')
    ax2.set_xlabel('Tiempo')
    ax2.set_title('Conservación')
    ax2.legend()
    ax2.grid(True)
    
    # 3. Hamiltoniano dependiente del tiempo (elementos seleccionados)
    ax3 = fig.add_subplot(gs[1, :])
    H_evolution = results['H_evolution']
    H_diagonal = [np.real(np.diag(H)) for H in H_evolution]
    H_diagonal = np.array(H_diagonal).T
    
    for i in range(min(3, H_diagonal.shape[0])):  # Máximo 3 elementos diagonales
        ax3.plot(time, H_diagonal[i], label=f'H_{{{i},{i}}}')
    
    # También graficar algunos elementos off-diagonal si existen
    if H_evolution[0].shape[0] > 1:
        H_offdiag = [np.real(H[0, 1]) for H in H_evolution]
        ax3.plot(time, H_offdiag, '--', label='Re(H_{0,1})', alpha=0.7)
    
    ax3.set_xlabel('Tiempo')
    ax3.set_ylabel('H(t) [elementos]')
    ax3.set_title('Evolución del Hamiltoniano')
    ax3.legend()
    ax3.grid(True)
    
    # 4. Otros observables (coherencias, etc.)
    ax4 = fig.add_subplot(gs[2, :2])
    for obs_name, obs_vals in observables.items():
        if 'Coherence' in obs_name or 'sigma' in obs_name:
            ax4.plot(time, np.real(obs_vals), label=f'Re({obs_name})')
            ax4.plot(time, np.imag(obs_vals), '--', label=f'Im({obs_name})', alpha=0.7)
    ax4.set_xlabel('Tiempo')
    ax4.set_ylabel('Coherencias')
    ax4.set_title('Evolución de Coherencias')
    ax4.legend()
    ax4.grid(True)
    
    # 5. Matriz de densidad final (representación)
    ax5 = fig.add_subplot(gs[2, 2])
    rho_final = results['rho_t'][-1]
    im = ax5.imshow(np.abs(rho_final), cmap='viridis', interpolation='nearest')
    ax5.set_title('|ρ(t_final)|')
    ax5.set_xlabel('j')
    ax5.set_ylabel('i')
    plt.colorbar(im, ax=ax5)
    
    plt.suptitle(f'Resultados: {results["name"]} (Dependencia Temporal)', fontsize=14)
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        logger.info(f"Figura guardada en {save_path}")
    
    plt.show()

def demonstrate_time_dependent_systems():
    """Función de demostración para sistemas con dependencia temporal."""
    
    print("\n" + "="*70)
    print("SIMULACIÓN DE SISTEMAS CON DEPENDENCIA TEMPORAL")
    print("="*70)
    
    simulator = TimeDependentLindladSimulator()
    
    # Ejemplo 1: Qubit controlado
    print("\n1. Simulando qubit con pulso láser...")
    driven_qubit = create_driven_qubit_system()
    results_qubit = simulator.simulate(driven_qubit)
    
    if results_qubit['success']:
        print("   ✓ Simulación exitosa")
        plot_time_dependent_results(results_qubit)
    
    # Ejemplo 2: STIRAP
    print("\n2. Simulando transferencia adiabática (STIRAP)...")
    stirap_system = create_adiabatic_passage_system()
    results_stirap = simulator.simulate(stirap_system)
    
    if results_stirap['success']:
        print("   ✓ Simulación exitosa")
        plot_time_dependent_results(results_stirap)
    
    # Ejemplo 3: Oscilador paramétrico
    print("\n3. Simulando oscilador paramétrico...")
    param_osc = create_parametric_oscillator()
    results_osc = simulator.simulate(param_osc)
    
    if results_osc['success']:
        print("   ✓ Simulación exitosa")
        plot_time_dependent_results(results_osc)

if __name__ == "__main__":
    demonstrate_time_dependent_systems()