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)


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
        
        # Encoder: time series → latent parameters
        encoder_list = []
        in_dim = input_dim
        for i in range(encoder_layers):
            out_dim = latent_dim if i == encoder_layers - 1 else input_dim
            encoder_list.extend([
                nn.Linear(in_dim, out_dim),
                nn.GELU() if i < encoder_layers - 1 else nn.Identity()
            ])
            in_dim = input_dim
            
        self.encoder = nn.Sequential(*encoder_list)
        
        # 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
        """
        # Global average pooling
        x_pooled = x.mean(dim=1)  # (batch, input_dim)
        
        # Encode to latent space
        z = self.encoder(x_pooled)  # (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
        """
        if t.dim() == 1:
            t = t.unsqueeze(0)
            
        freq = params['frequencies']    # (batch, n_basis)
        amp = params['amplitudes']      # (batch, n_basis)
        phase = params['phases']        # (batch, n_basis)
        
        # Compute Fourier series
        t_expanded = t.unsqueeze(-1)    # (batch, seq, 1)
        freq_expanded = freq.unsqueeze(1)  # (batch, 1, n_basis)
        amp_expanded = amp.unsqueeze(1)     # (batch, 1, n_basis)
        phase_expanded = phase.unsqueeze(1)  # (batch, 1, n_basis)
        
        fourier_sum = (amp_expanded * torch.sin(
            freq_expanded * t_expanded + phase_expanded
        )).sum(dim=-1)  # (batch, seq)
        
        return fourier_sum
    
    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,
    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
        verbose: Print progress
        
    Returns:
        Training history dictionary
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
    
    history = {'loss': [], 'reconstruction_error': [], 'param_cost': []}
    
    t = torch.linspace(0, 1, time_series.shape[1], device=time_series.device)
    if time_series.dim() > 2:
        t = t.unsqueeze(0)
    
    for epoch in range(epochs):
        optimizer.zero_grad()
        
        # Forward pass
        reconstruction, params = model(time_series, t)
        
        # Reconstruction loss (collapse uncertainty)
        recon_loss = F.mse_loss(reconstruction, time_series[..., 0])  # Assume 1D for simplicity
        
        # Parameter cost (regularization)
        param_cost = sum(p.abs().sum() for p in model.parameters())
        
        # Total loss: balance reconstruction vs compression
        loss = recon_loss + lambda_reg * param_cost
        
        # Backward pass
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        
        # Record history
        history['loss'].append(loss.item())
        history['reconstruction_error'].append(recon_loss.item())
        history['param_cost'].append(param_cost.item())
        
        if verbose and epoch % 100 == 0:
            print(f"Epoch {epoch}: Loss={loss.item():.4f}, "
                  f"Recon={recon_loss.item():.4f}, "
                  f"Params={param_cost.item():.4f}")
    
    return history


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

# 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)")