#!/usr/bin/env python3
"""
encoder.py - Layer Encoder using SVD
====================================

Encodes tensors using truncated SVD for compression.
"""

import hashlib
from typing import Dict, Any, Optional
import numpy as np


class LayerEncoder:
    """Encode extracted tensors using truncated SVD."""

    def __init__(self, k_effective: int = 64, n_subvectors: int = 8, n_centroids: int = 256):
        self.k_effective = k_effective
        self.n_subvectors = n_subvectors
        self.n_centroids = n_centroids

    def encode(self, tensor: np.ndarray, name: str) -> Dict[str, Any]:
        """Encode tensor using truncated SVD + statistics."""
        M, N = tensor.shape if tensor.ndim == 2 else (tensor.shape[0], np.prod(tensor.shape[1:]))
        tensor_2d = tensor.reshape(M, -1).astype(np.float32)

        k = min(self.k_effective, min(M, N) - 1)

        # Compute SVD
        U_k, S_k, Vt_k = None, None, None
        recon_error = float('inf')
        energy_ratio = 0.0

        try:
            U, S, Vt = np.linalg.svd(tensor_2d, full_matrices=False)
            U_k = U[:, :k]
            S_k = S[:k]
            Vt_k = Vt[:k, :]

            # Reconstruction error
            W_approx = U_k @ np.diag(S_k) @ Vt_k
            recon_error = float(np.max(np.abs(tensor_2d - W_approx)))

            # Energy retained
            total_energy = float(np.sum(S ** 2))
            retained_energy = float(np.sum(S_k ** 2))
            energy_ratio = retained_energy / total_energy if total_energy > 0 else 0

        except np.linalg.LinAlgError:
            pass

        # Compute structural statistics
        stats = {
            'mean': float(np.mean(tensor)),
            'std': float(np.std(tensor)),
            'min': float(np.min(tensor)),
            'max': float(np.max(tensor)),
            'norm': float(np.linalg.norm(tensor)),
            'sparsity': float(np.sum(np.abs(tensor) < 1e-6) / tensor.size),
        }

        return {
            'name': name,
            'original_shape': list(tensor.shape),
            'k_effective': k,
            'recon_error': recon_error,
            'energy_ratio': energy_ratio,
            'statistics': stats,
            'sha256': hashlib.sha256(tensor.tobytes()).hexdigest(),
            '_U_k': U_k,
            '_S_k': S_k,
            '_Vt_k': Vt_k,
        }

    def reconstruct(self, encoded: Dict[str, Any]) -> Optional[np.ndarray]:
        """Reconstruct tensor from SVD components."""
        U_k = encoded.get('_U_k')
        S_k = encoded.get('_S_k')
        Vt_k = encoded.get('_Vt_k')

        if U_k is None or S_k is None or Vt_k is None:
            return None

        return U_k @ np.diag(S_k) @ Vt_k

    def save_encoded(self, encoded: Dict[str, Any], path: str) -> int:
        """Save encoded tensor to .npz file. Returns file size."""
        from pathlib import Path

        save_dict = {
            'meta': np.array([str({k: v for k, v in encoded.items() if not k.startswith('_')})], dtype='U10000'),
        }

        if encoded['_U_k'] is not None:
            save_dict['U_k'] = encoded['_U_k']
            save_dict['S_k'] = encoded['_S_k']
            save_dict['Vt_k'] = encoded['_Vt_k']

        np.savez_compressed(path, **save_dict)
        return Path(path).stat().st_size

    @staticmethod
    def load_encoded(path: str) -> Dict[str, Any]:
        """Load encoded tensor from .npz file."""
        data = np.load(path, allow_pickle=True)

        result = {}
        if 'meta' in data:
            import ast
            meta_str = str(data['meta'][0])
            try:
                result = ast.literal_eval(meta_str)
            except:
                result = {}

        if 'U_k' in data:
            result['_U_k'] = data['U_k']
            result['_S_k'] = data['S_k']
            result['_Vt_k'] = data['Vt_k']

        return result
