import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Tuple, Optional, Dict, List

class ParametricFourierBasis(nn.Module):
    """
    Learnable parametric Fourier basis functions.
    Instead of storing coefficients explicitly, we learn parameters
    that generate the Fourier coefficients on demand.
    
    F(t) = Σ θ_k * φ_k(t) where φ_k are basis functions
    """
    
    def __init__(
        self, 
        n_harmonics: int = 10,
        learnable_frequencies: bool = True,
        initial_freq_scale: float = 1.0
    ):
        super().__init__()
        self.n_harmonics = n_harmonics
        
        # Learnable frequencies (parameterized, not explicit)
        self.freq_base = nn.Parameter(
            torch.ones(n_harmonics) * initial_freq_scale
        )
        self.freq_multiplier = nn.Parameter(
            torch.arange(1, n_harmonics + 1).float()
        )
        
        # Learnable amplitudes (parameterized)
        self.amp_base = nn.Parameter(torch.randn(n_harmonics) * 0.1)
        self.amp_scale = nn.Parameter(torch.ones(n_harmonics))
        
        # Learnable phases
        self.phase_base = nn.Parameter(
            torch.rand(n_harmonics) * 2 * math.pi
        )
        
    def get_frequencies(self) -> torch.Tensor:
        """Get frequencies from parameters - O(n) not infinite."""
        return self.freq_base * self.freq_multiplier
    
    def get_amplitudes(self) -> torch.Tensor:
        """Get amplitudes from parameters - O(n) not infinite."""
        return F.softplus(self.amp_base) * self.amp_scale
    
    def get_phases(self) -> torch.Tensor:
        """Get phases from parameters."""
        return torch.tanh(self.phase_base) * math.pi
    
    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """
        Generate Fourier series value at time t.
        
        Args:
            t: Time points, shape (batch, seq_len) or (seq_len,)
            
        Returns:
            Fourier series values, same shape as t
        """
        if t.dim() == 1:
            t = t.unsqueeze(0)
            
        batch_size = t.shape[0]
        seq_len = t.shape[1]
        
        # Get parametric coefficients
        omega = self.get_frequencies()  # (n_harmonics,)
        A = self.get_amplitudes()       # (n_harmonics,)
        phi = self.get_phases()         # (n_harmonics,)
        
        # Compute Fourier basis: A_k * sin(k*omega*t + phi_k)
        # Shape: (batch, seq, n_harmonics)
        t_expanded = t.unsqueeze(-1)    # (batch, seq, 1)
        harmonics = A * torch.sin(
            omega * t_expanded + phi
        )  # (batch, seq, n_harmonics)
        
        # Sum over harmonics
        return harmonics.sum(dim=-1)    # (batch, seq)

    def inverse(
        self,
        frequencies: torch.Tensor,
        amplitudes: torch.Tensor,
        phases: torch.Tensor,
        t: torch.Tensor
    ) -> torch.Tensor:
        """
        Reconstruct a time-domain signal from explicit Fourier parameters.

        This is the inverse transform for this parameterization, similar in
        spirit to an inverse FFT, but using learned continuous frequencies
        instead of discrete FFT bins.

        Args:
            frequencies: Frequency parameters, shape (n_harmonics,) or (batch, n_harmonics)
            amplitudes: Amplitude parameters, same shape as frequencies
            phases: Phase parameters, same shape as frequencies
            t: Time points, shape (seq_len,) or (batch, seq_len)

        Returns:
            Reconstructed time-domain signal, shape (batch, seq_len)
        """
        if t.dim() == 1:
            t = t.unsqueeze(0)

        if frequencies.dim() == 1:
            frequencies = frequencies.unsqueeze(0)
        if amplitudes.dim() == 1:
            amplitudes = amplitudes.unsqueeze(0)
        if phases.dim() == 1:
            phases = phases.unsqueeze(0)

        if frequencies.shape[0] == 1 and t.shape[0] > 1:
            frequencies = frequencies.expand(t.shape[0], -1)
            amplitudes = amplitudes.expand(t.shape[0], -1)
            phases = phases.expand(t.shape[0], -1)

        t_expanded = t.unsqueeze(-1)                  # (batch, seq, 1)
        freq_expanded = frequencies.unsqueeze(1)      # (batch, 1, n_harmonics)
        amp_expanded = amplitudes.unsqueeze(1)        # (batch, 1, n_harmonics)
        phase_expanded = phases.unsqueeze(1)          # (batch, 1, n_harmonics)

        return (amp_expanded * torch.sin(
            freq_expanded * t_expanded + phase_expanded
        )).sum(dim=-1)


class LowRankFourierLayer(nn.Module):
    """
    Low-rank parametric Fourier layer.
    W_ij = Σ_k u_ik * v_jk (rank-r factorization of weight matrix)
    
    Compression: O(mn) → O(r(m+n))
    """
    
    def __init__(self, input_dim: int, output_dim: int, rank: int = 5):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.rank = rank
        
        # Factorized parameters: U (m×r) and V (n×r)
        # Instead of storing W (m×n), we store U and V
        self.U = nn.Parameter(torch.randn(output_dim, rank) * 0.01)
        self.V = nn.Parameter(torch.randn(input_dim, rank) * 0.01)
        
        # Optional: learnable basis scaling
        self.basis_scale = nn.Parameter(torch.ones(rank))
        
    def get_weight_matrix(self) -> torch.Tensor:
        """
        Reconstruct full weight matrix from parametric factors.
        Only computed when needed (on-demand generation).
        """
        # W ≈ U @ V^T with element-wise scaling
        W = self.U @ self.V.T  # (output_dim, input_dim)
        return W * self.basis_scale.unsqueeze(0)
    
    def forward(
        self, 
        x: torch.Tensor, 
        generate_weights: bool = False
    ) -> torch.Tensor:
        """
        Forward pass with optional on-demand weight generation.
        
        Args:
            x: Input tensor (batch, seq, input_dim)
            generate_weights: If True, reconstructs full matrix
                            If False, uses factorized form directly
        """
        if generate_weights:
            W = self.get_weight_matrix()  # (output_dim, input_dim)
            return torch.matmul(x, W.T)
        else:
            # Direct factorized computation: x @ (V @ U^T)^T = x @ V @ U^T
            # More memory efficient, avoids full matrix
            return torch.matmul(
                torch.matmul(x, self.V),  # (batch, seq, rank)
                self.U.T                  # (rank, output_dim)
            )  # (batch, seq, output_dim)
    
    @property
    def compression_ratio(self) -> float:
        """How much this layer is compressed vs explicit storage."""
        explicit = self.input_dim * self.output_dim
        parametric = self.rank * (self.input_dim + self.output_dim) + self.rank
        return explicit / parametric


class FourierGeneratorLayer(nn.Module):
    """
    Neural network generator that produces Fourier-based weights.
    
    W_ij = NN_θ(i, j) — a small network generates weight values
    from their indices in the weight matrix.
    
    Compression: O(mn) → O(θ) where θ << mn
    """
    
    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dim: int = 32,
        n_fourier_terms: int = 8
    ):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.n_fourier_terms = n_fourier_terms
        
        # Fourier basis parameters (stationary component)
        self.register_buffer(
            'freq_i', 
            torch.arange(1, n_fourier_terms + 1).float()
        )
        self.register_buffer(
            'freq_j', 
            torch.arange(1, n_fourier_terms + 1).float()
        )
        
        # Learnable coefficients for Fourier basis (probability component)
        # Shape: (n_fourier_terms, n_fourier_terms)
        self.fourier_coeffs = nn.Parameter(
            torch.randn(n_fourier_terms, n_fourier_terms) * 0.01
        )
        
        # Optional neural correction layer
        self.neural_correction = nn.Sequential(
            nn.Linear(2, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1)
        )
        
    def forward(self, indices_i: torch.Tensor, indices_j: torch.Tensor) -> torch.Tensor:
        """
        Generate weight values for given matrix indices.
        
        Args:
            indices_i: Row indices (output dimension)
            indices_j: Column indices (input dimension)
            
        Returns:
            Weight values for each (i,j) pair
        """
        # Normalize indices to [-1, 1]
        i_norm = 2 * indices_i.float() / self.output_dim - 1
        j_norm = 2 * indices_j.float() / self.input_dim - 1
        
        # Fourier basis component: Σ c_kl * sin(ki * ω) * cos(lj * ω)
        i_expanded = i_norm.unsqueeze(-1)  # (..., 1)
        j_expanded = j_norm.unsqueeze(-1)  # (..., 1)
        
        # Compute Fourier features
        i_freq = torch.sin(
            i_expanded * self.freq_i.unsqueeze(0) * math.pi
        )  # (..., n_fourier_terms)
        j_freq = torch.cos(
            j_expanded * self.freq_j.unsqueeze(0) * math.pi
        )  # (..., n_fourier_terms)
        
        # Bilinear Fourier interaction for each index pair: (..., n) x (n, n) x (..., n) -> (...)
        fourier_out = torch.einsum("...k,kl,...l->...", i_freq, self.fourier_coeffs, j_freq)
        
        # Neural correction for residual
        index_features = torch.stack([i_norm, j_norm], dim=-1)
        neural_correction = self.neural_correction(index_features).squeeze(-1)
        
        return fourier_out + neural_correction


class ParametricTimeSeriesModel(nn.Module):
    """
    Complete parametric time series model using learned Fourier series.
    
    Instead of storing explicit weights, uses parametric functions
    to generate weights on demand.
    """
    
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        output_dim: int,
        n_layers: int = 3,
        rank: int = 8,
        n_fourier_terms: int = 10
    ):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.n_layers = n_layers
        self.input_with_fourier_dim = input_dim + 1
        
        # Input embedding with parametric Fourier basis
        self.input_basis = ParametricFourierBasis(n_fourier_terms)
        self.input_projection = nn.Linear(self.input_with_fourier_dim, hidden_dim)
        
        # Hidden layers using low-rank factorization
        self.hidden_layers = nn.ModuleList([
            LowRankFourierLayer(
                hidden_dim,
                hidden_dim,
                rank=rank
            )
            for i in range(n_layers)
        ])
        
        # Output projection with Fourier generator
        self.output_projection = FourierGeneratorLayer(
            hidden_dim, output_dim, n_fourier_terms=n_fourier_terms
        )
        
        # ODE-like transition parameters
        self.ode_gain = nn.Parameter(torch.ones(n_layers))
        self.ode_damping = nn.Parameter(torch.ones(n_layers) * 0.9)
        
    def forward(
        self, 
        x: torch.Tensor,
        return_generated_weights: bool = False
    ) -> torch.Tensor:
        """
        Forward pass with parametric weight generation.
        
        Args:
            x: Input time series (batch, seq_len, input_dim)
            return_generated_weights: If True, return generated weights too
            
        Returns:
            Output predictions and optionally generated weights
        """
        batch_size, seq_len, _ = x.shape
        
        # Embed input using parametric Fourier basis
        t = torch.linspace(0, 1, seq_len, device=x.device)
        fourier_features = self.input_basis(t).repeat(batch_size, 1)  # (batch, seq)
        
        # Concatenate input with Fourier features
        h = torch.cat([x, fourier_features.unsqueeze(-1)], dim=-1)
        h = self.input_projection(h)
        
        # Hidden layers with ODE-style updates
        for layer_idx, layer in enumerate(self.hidden_layers):
            # Generate weights on demand (parametric)
            h_new = layer(h, generate_weights=False)  # Uses factorized form
            
            # ODE update: h = h + gain * (new_h - damping * h)
            h = h + self.ode_gain[layer_idx] * (
                h_new - self.ode_damping[layer_idx] * h
            )
            
            h = F.gelu(h)  # Non-linearity
            
        # Output projection using Fourier generator
        seq_idx = torch.arange(seq_len, device=x.device).view(1, seq_len, 1).expand(
            batch_size, seq_len, self.hidden_dim
        )
        hidden_idx = torch.arange(self.hidden_dim, device=x.device).view(1, 1, self.hidden_dim).expand(
            batch_size, seq_len, self.hidden_dim
        )
        
        # Generate output weights parametrically for each hidden feature
        output_weights = self.output_projection(seq_idx, hidden_idx)  # (batch, seq, hidden)
        
        # Apply generated weights to hidden states
        output = (h * output_weights).sum(dim=-1, keepdim=True)  # (batch, seq, 1)
        
        if return_generated_weights:
            return output, output_weights
        return output
    
    def get_storage_cost(self) -> Dict[str, int]:
        """Calculate storage cost of parametric vs explicit representation."""
        explicit_params = self.input_dim * self.hidden_dim * self.n_layers
        
        parametric_params = sum(
            layer.U.numel() + layer.V.numel() + layer.basis_scale.numel()
            for layer in self.hidden_layers
        )
        parametric_params += self.input_basis.n_harmonics * 6  # frequencies, amps, phases
        parametric_params += self.output_projection.fourier_coeffs.numel()
        parametric_params += self.output_projection.neural_correction.state_dict().__len__()
        
        return {
            'explicit': explicit_params,
            'parametric': parametric_params,
            'compression_ratio': explicit_params / max(parametric_params, 1)
        }


class AdaptiveFourierEncoder(nn.Module):
    """
    Encodes arbitrary time series into parametric Fourier representation.
    
    Given a time series x(t), finds parameters θ such that:
    x(t) ≈ Σ_k θ_k * φ_k(t)
    
    The encoder learns to compress arbitrary inputs into fixed-size parameter vectors.
    """
    
    def __init__(
        self,
        input_dim: int,
        latent_dim: int = 32,
        n_fourier_basis: int = 16,
        encoder_layers: int = 3
    ):
        super().__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.n_fourier_basis = n_fourier_basis
        
        # Time-aware encoder: preserve ordering instead of collapsing with mean pooling.
        self.encoder = nn.GRU(
            input_size=input_dim,
            hidden_size=latent_dim,
            num_layers=encoder_layers,
            batch_first=True
        )
        
        # Parametric Fourier decoder
        self.fourier_basis = ParametricFourierBasis(n_fourier_basis)
        
        # Latent to frequency/amplitude/phase mapping
        self.latent_to_freq = nn.Linear(latent_dim, n_fourier_basis)
        self.latent_to_amp = nn.Linear(latent_dim, n_fourier_basis)
        self.latent_to_phase = nn.Linear(latent_dim, n_fourier_basis)
        
    def encode(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Encode time series into parametric representation.
        
        Args:
            x: Time series (batch, seq_len, input_dim)
            
        Returns:
            Dictionary of parametric components
        """
        # Encode the full sequence so temporal order remains available to the latent code.
        _, h_n = self.encoder(x)  # h_n: (num_layers, batch, latent_dim)
        z = h_n[-1]               # (batch, latent_dim)
        
        # Map to Fourier parameters
        return {
            'latent': z,
            'frequencies': F.softplus(self.latent_to_freq(z)),
            'amplitudes': torch.tanh(self.latent_to_amp(z)),
            'phases': torch.tanh(self.latent_to_phase(z)) * math.pi
        }
    
    def decode(self, params: Dict[str, torch.Tensor], t: torch.Tensor) -> torch.Tensor:
        """
        Decode parametric representation back to time series.
        
        Args:
            params: Dictionary from encode()
            t: Time points (seq_len,) or (batch, seq_len)
            
        Returns:
            Reconstructed time series
        """
        return self.fourier_basis.inverse(
            params["frequencies"],
            params["amplitudes"],
            params["phases"],
            t
        )

    def inverse(self, params: Dict[str, torch.Tensor], t: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Alias for decode(), matching the inverse-transform naming you asked for.

        If `t` is omitted, a default evenly spaced grid over [0, 1] is used.
        """
        if t is None:
            batch = params["frequencies"].shape[0] if params["frequencies"].dim() > 1 else 1
            n_steps = getattr(self, "_default_inverse_steps", 100)
            t = torch.linspace(0, 1, n_steps, device=params["frequencies"].device)
            if batch > 1:
                t = t.unsqueeze(0).expand(batch, -1)
        return self.decode(params, t)
    
    def forward(self, x: torch.Tensor, t: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Dict]:
        """
        Full encode-decode cycle.
        
        Args:
            x: Input time series
            t: Time points (if None, uses linspace)
            
        Returns:
            Reconstructed time series and parameters
        """
        if t is None:
            t = torch.linspace(0, 1, x.shape[1], device=x.device)
            if x.dim() > 2:
                t = t.unsqueeze(0)
                
        params = self.encode(x)
        reconstruction = self.decode(params, t)
        
        return reconstruction, params


def train_parametric_fourier(
    model: AdaptiveFourierEncoder,
    time_series: torch.Tensor,
    epochs: int = 1000,
    lr: float = 1e-3,
    lambda_reg: float = 1e-4,
    lambda_freq_reg: float = 1e-4,
    lambda_amp_reg: float = 1e-4,
    lambda_phase_reg: float = 1e-4,
    grad_clip: float = 1.0,
    patience: Optional[int] = None,
    min_delta: float = 1e-6,
    use_inverse: bool = True,
    scheduler_patience: Optional[int] = None,
    scheduler_factor: float = 0.8,
    scheduler_min_lr: float = 1e-5,
    scheduler_on: str = "reconstruction",
    verbose: bool = True
) -> Dict[str, List]:
    """
    Train the parametric Fourier encoder on a time series.
    
    This implements the retro-compression idea: given explicit data,
    find the parametric function that best approximates it.
    
    Args:
        model: AdaptiveFourierEncoder to train
        time_series: Training data (batch, seq_len, input_dim)
        epochs: Number of training epochs
        lr: Learning rate
        lambda_reg: L2 regularization strength
        lambda_freq_reg: Penalty on learned frequencies
        lambda_amp_reg: Penalty on learned amplitudes
        lambda_phase_reg: Penalty on learned phases
        grad_clip: Gradient clipping threshold
        patience: Early stopping patience. Disable if None.
        min_delta: Minimum improvement required to reset patience
        use_inverse: If True, train through model.inverse(params, t)
        scheduler_patience: Plateau scheduler patience. Defaults to a conservative value.
        scheduler_factor: LR decay factor for plateau scheduler
        scheduler_min_lr: Minimum learning rate
        scheduler_on: Metric used for LR scheduling: "loss" or "reconstruction"
        verbose: Print progress
        
    Returns:
        Training history dictionary
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    if scheduler_patience is None:
        scheduler_patience = max(20, epochs // 10)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=scheduler_factor,
        patience=scheduler_patience,
        min_lr=scheduler_min_lr
    )
    
    history = {
        'loss': [],
        'reconstruction_error': [],
        'param_cost': [],
        'frequency_cost': [],
        'amplitude_cost': [],
        'phase_cost': [],
        'lr': [],
    }
    
    t = torch.linspace(0, 1, time_series.shape[1], device=time_series.device)
    if time_series.dim() > 2:
        t = t.unsqueeze(0)

    target = time_series[..., 0] if time_series.shape[-1] == 1 else time_series.squeeze(-1)

    best_loss = float("inf")
    best_state = None
    bad_epochs = 0
    
    for epoch in range(epochs):
        optimizer.zero_grad()
        
        # Forward pass through the encoder, then decode explicitly via the inverse path.
        _, params = model(time_series, t)
        reconstruction = model.inverse(params, t) if use_inverse else model.decode(params, t)
        if reconstruction.dim() == 3 and reconstruction.shape[-1] == 1:
            reconstruction = reconstruction.squeeze(-1)
        
        # Reconstruction loss on the time domain.
        recon_loss = F.mse_loss(reconstruction, target)
        
        # Parameter regularization, normalized so it doesn't explode with model size.
        param_cost = sum(p.square().mean() for p in model.parameters())
        frequency_cost = params["frequencies"].square().mean()
        amplitude_cost = params["amplitudes"].square().mean()
        phase_cost = params["phases"].square().mean()
        
        # Total loss balances reconstruction fidelity and parameter smoothness.
        loss = (
            recon_loss
            + lambda_reg * param_cost
            + lambda_freq_reg * frequency_cost
            + lambda_amp_reg * amplitude_cost
            + lambda_phase_reg * phase_cost
        )
        
        # Backward pass
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
        optimizer.step()
        scheduler_metric = recon_loss.item() if scheduler_on == "reconstruction" else loss.item()
        scheduler.step(scheduler_metric)
        
        # Record history
        history['loss'].append(loss.item())
        history['reconstruction_error'].append(recon_loss.item())
        history['param_cost'].append(param_cost.item())
        history['frequency_cost'].append(frequency_cost.item())
        history['amplitude_cost'].append(amplitude_cost.item())
        history['phase_cost'].append(phase_cost.item())
        history['lr'].append(optimizer.param_groups[0]['lr'])

        if loss.item() + min_delta < best_loss:
            best_loss = loss.item()
            best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
            bad_epochs = 0
        else:
            bad_epochs += 1
        
        if verbose and (epoch % 100 == 0 or epoch == epochs - 1):
            print(
                f"Epoch {epoch}: Loss={loss.item():.4f}, "
                f"Recon={recon_loss.item():.4f}, "
                f"Param={param_cost.item():.4f}, "
                f"Freq={frequency_cost.item():.4f}, "
                f"Amp={amplitude_cost.item():.4f}, "
                f"Phase={phase_cost.item():.4f}, "
                f"Sched={scheduler_metric:.4f}, "
                f"LR={optimizer.param_groups[0]['lr']:.2e}"
            )

        if patience is not None and bad_epochs >= patience:
            if verbose:
                print(f"Early stopping at epoch {epoch} (best loss {best_loss:.4f}).")
            break

    if best_state is not None:
        model.load_state_dict(best_state)
    
    return history


# =============================================================================
# EXAMPLE USAGE
# =============================================================================

if __name__ == "__main__":
    # Create a synthetic time series
    torch.manual_seed(42)
    
    seq_len = 10000
    batch_size = 32
    input_dim = 1
    
    # Generate multi-frequency time series
    t = torch.linspace(0, 4 * torch.pi, seq_len)
    x = (
        2.0 * torch.sin(t) +           # Low frequency
        0.5 * torch.sin(5 * t) +       # Medium frequency
        0.25 * torch.sin(10 * t) +     # High frequency
        0.1 * torch.randn(seq_len)     # Noise
    )

    #x = torch.rand(seq_len)
    
    # Create batch
    time_series = x.unsqueeze(0).unsqueeze(-1).expand(batch_size, -1, -1)
    
    print("=" * 60)
    print("PARAMETRIC FOURIER SERIES - LEARNED FROM TIME SERIES")
    print("=" * 60)
    
    # Initialize model
    model = AdaptiveFourierEncoder(
        input_dim=input_dim,
        latent_dim=16,
        n_fourier_basis=8,
        encoder_layers=3
    )
    
    # Train
    print("\n[Training parametric model...]")
    history = train_parametric_fourier(
        model, 
        time_series,
        epochs=500,
        lr=1e-2,
        lambda_reg=1e-4,
        verbose=True
    )
    
    # Evaluate
    print("\n[Evaluating...]")
    with torch.no_grad():
        reconstruction, params = model(time_series[:1])
    
    # Calculate compression
    explicit_params = seq_len * input_dim
    parametric_params = sum(p.numel() for p in model.parameters())
    
    print(f"\n{'='*60}")
    print("COMPRESSION RESULTS:")
    print(f"  Explicit storage: {explicit_params} values")
    print(f"  Parametric storage: {parametric_params} parameters")
    print(f"  Compression ratio: {explicit_params/parametric_params:.1f}x")
    
    print(f"\nLearned Fourier parameters:")
    print(f"  Frequencies: {params['frequencies'][0].round(decimals=3)}")
    print(f"  Amplitudes:  {params['amplitudes'][0].round(decimals=3)}")
    print(f"  Phases:      {params['phases'][0].round(decimals=3)}")
    
    # Test the full parametric time series model
    print(f"\n{'='*60}")
    print("FULL PARAMETRIC TIME SERIES MODEL:")
    print("=" * 60)
    
    full_model = ParametricTimeSeriesModel(
        input_dim=input_dim,
        hidden_dim=32,
        output_dim=1,
        n_layers=3,
        rank=5,
        n_fourier_terms=10
    )
    
    # Forward pass
    output = full_model(time_series)
    print(f"Input shape: {time_series.shape}")
    print(f"Output shape: {output.shape}")
    
    # Storage cost analysis
    storage = full_model.get_storage_cost()
    print(f"\nStorage analysis:")
    print(f"  Explicit parameters: {storage['explicit']}")
    print(f"  Parametric parameters: {storage['parametric']}")
    print(f"  Compression ratio: {storage['compression_ratio']:.1f}x")
    
    # Generate weights on demand
    print(f"\n[Testing on-demand weight generation...]")
    with torch.no_grad():
        output, weights = full_model(time_series, return_generated_weights=True)
    print(f"Generated weights shape: {weights.shape}")
    print(f"(Weights generated on-the-fly, not stored)")