"""
CDNA v3 tensor writer.

Writes per-tensor directories with the v3 layout:
  {safe_name}.cdnav3/
    meta.json       - shape, name, policy, codebook hash
    codebook.npy    - [N] float32 codebook
    indices.bin     - [M*K] uint8 raw indices
    sidecar.npz     - outlier positions + values (optional)
    stats.json      - encoding fidelity metrics
"""

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import numpy as np

from helix_substrate.tensor_policy import (
    TensorClass,
    TensorPolicy,
    classify_tensor,
    get_default_policy,
)
from helix_substrate.generate_sidecars_v3 import (
    find_outliers_percentile,
    write_sidecar_npz,
)
# Lazy import: morpho_codec requires scipy (heavy dep, dead codec)
morpho_encode = None

# Max elements to subsample for k-means (prevents OOM)
_KMEANS_MAX_SAMPLES = 500_000
# Chunk size for assignment pass
_ASSIGN_CHUNK_SIZE = 1_024 * 1_024


class CDNAv3Writer:
    """Write tensors in CDNA v3 directory format."""

    def __init__(self, base_dir: Path):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def write_tensor(
        self,
        tensor: np.ndarray,
        tensor_name: str,
        policy: Optional[TensorPolicy] = None,
        source_artifact: str = "",
    ) -> dict:
        """
        Write a tensor in v3 format.

        Args:
            tensor: Input tensor (1D or 2D, float32)
            tensor_name: Canonical tensor name
            policy: Compression policy (auto-assigned if None)
            source_artifact: Source file path for provenance

        Returns:
            Stats dict with encoding metrics
        """
        if policy is None:
            tc = classify_tensor(tensor_name, shape=tensor.shape)
            policy = get_default_policy(tc)

        # Morpho codec: grow from seed (handles 1D and 2D)
        if policy.storage_mode == "morpho":
            tensor = tensor.astype(np.float32)
            return self._write_morpho(tensor, tensor_name, policy)

        # 1D tensors and exact-policy tensors: save as float16 (not float32)
        # to match typical BF16/FP16 source models and avoid 2x size inflation.
        if tensor.ndim == 1 or policy.storage_mode == "exact":
            return self._write_exact(tensor.astype(np.float16), tensor_name)

        # Cast to float32 for k-means quantization (needs FP32 precision)
        tensor = tensor.astype(np.float32)

        # 2D tensors: quantize with optional sidecar
        if tensor.ndim > 2:
            tensor = tensor.reshape(-1, tensor.shape[-1])

        return self._write_quantized(tensor, tensor_name, policy, source_artifact)

    def _write_exact(self, tensor: np.ndarray, tensor_name: str) -> dict:
        """Write a tensor exactly as .npy with a companion meta file."""
        safe = _safe_name(tensor_name)
        out_path = self.base_dir / f"{safe}.npy"
        np.save(out_path, tensor)

        # Write companion meta so manifest can recover the original name
        meta_path = self.base_dir / f"{safe}.npy.meta.json"
        meta_path.write_text(json.dumps({
            "tensor_name": tensor_name,
            "shape": list(tensor.shape),
            "storage_mode": "exact",
        }))

        return {
            "tensor_name": tensor_name,
            "storage_mode": "exact",
            "shape": list(tensor.shape),
            "original_bytes": tensor.nbytes,
            "compressed_bytes": out_path.stat().st_size,
        }

    def _write_morpho(
        self,
        tensor: np.ndarray,
        tensor_name: str,
        policy: TensorPolicy,
    ) -> dict:
        """Write a tensor using morpho codec (wave-grown from seed).
        Requires scipy: pip install scipy

        If policy.morpho_min_cosine > 0, applies a quality gate:
        falls back to exact storage if the morpho fit doesn't meet threshold.
        """
        global morpho_encode
        if morpho_encode is None:
            from helix_substrate.morpho_codec import morpho_encode as _me
            morpho_encode = _me

        safe = _safe_name(tensor_name)
        out_dir = self.base_dir / f"{safe}.cdnav3"

        stats = morpho_encode(
            tensor=tensor,
            tensor_name=tensor_name,
            output_dir=out_dir,
            steps=policy.morpho_growth_steps,
            target_se=policy.morpho_target_se,
            c=policy.morpho_c,
            gamma=policy.morpho_gamma,
            geometry=policy.morpho_geometry,
            fit=policy.morpho_fit,
            n_codons=policy.morpho_n_codons,
            max_iter=policy.morpho_max_iter,
            verbose=policy.morpho_verbose,
            spectral=policy.morpho_spectral,
            n_harmonics=policy.morpho_n_harmonics,
        )

        # Quality gate: fall back to exact if below threshold
        if policy.morpho_min_cosine > 0 and stats.get("cosine", 0) < policy.morpho_min_cosine:
            import shutil
            if out_dir.exists():
                shutil.rmtree(out_dir)
            exact_stats = self._write_exact(tensor, tensor_name)
            exact_stats["morpho_fallback"] = True
            exact_stats["morpho_cosine"] = stats.get("cosine", 0)
            return exact_stats

        return stats

    def _write_quantized(
        self,
        tensor: np.ndarray,
        tensor_name: str,
        policy: TensorPolicy,
        source_artifact: str,
    ) -> dict:
        """Write a quantized tensor with optional sidecar."""
        safe = _safe_name(tensor_name)
        out_dir = self.base_dir / f"{safe}.cdnav3"
        out_dir.mkdir(parents=True, exist_ok=True)

        rows, cols = tensor.shape
        flat = tensor.ravel()
        vector_dim = getattr(policy, "vector_dim", 1)

        # --- Build codebook + assign indices ---
        if vector_dim > 1:
            # Multi-dimensional VQ path
            assert cols % vector_dim == 0, (
                f"in_features ({cols}) must be divisible by vector_dim ({vector_dim})"
            )
            codebook, indices = self._build_vector_codebook_and_assign(
                tensor, policy.n_clusters, vector_dim
            )
            index_dtype = np.uint16 if policy.n_clusters > 256 else np.uint8
            indices = indices.astype(index_dtype)
            # indices shape: [rows * cols // vector_dim]
            indices_2d = indices.reshape(rows, cols // vector_dim)
        else:
            # Scalar VQ path (production default)
            if policy.use_kmeans:
                codebook = self._build_kmeans_codebook(flat, policy.n_clusters)
            else:
                codebook = self._build_uniform_codebook(flat, policy.n_clusters)
            index_dtype = np.uint16 if len(codebook) > 256 else np.uint8
            indices = self._chunked_assign(flat, codebook, index_dtype)
            indices_2d = indices.reshape(rows, cols)

        # --- Compute baseline fidelity (before sidecar) ---
        if vector_dim > 1:
            # Reconstruct: codebook[indices] → [N_vectors, d] → flatten to [M*N]
            reconstructed = codebook[indices.ravel()].reshape(-1)
        else:
            reconstructed = codebook[indices]
        cosine_no_sidecar = float(_cosine(flat, reconstructed))
        max_abs_diff = float(np.max(np.abs(flat - reconstructed)))

        # --- Sidecar ---
        num_outliers = 0
        sidecar_bytes = 0
        cosine_with_sidecar = cosine_no_sidecar

        if policy.sidecar_enabled:
            positions, values = find_outliers_percentile(tensor, policy.percentile)
            if len(positions) > 0:
                if policy.max_corrections > 0 and len(positions) > policy.max_corrections:
                    # Keep only top-N by absolute value
                    top_idx = np.argsort(np.abs(values))[::-1][:policy.max_corrections]
                    positions = np.sort(positions[top_idx])
                    values = tensor.ravel()[positions]

                sidecar_path = out_dir / "sidecar.npz"
                receipt = write_sidecar_npz(positions, values, sidecar_path)
                num_outliers = receipt["num_corrections"]
                sidecar_bytes = receipt["size_bytes"]

                # Compute fidelity with sidecar
                patched = reconstructed.copy()
                patched[positions] = values
                cosine_with_sidecar = float(_cosine(flat, patched))

        # --- Write codebook ---
        codebook_path = out_dir / "codebook.npy"
        np.save(codebook_path, codebook)
        codebook_sha256 = hashlib.sha256(codebook.tobytes()).hexdigest()

        # --- Write indices ---
        indices_path = out_dir / "indices.bin"
        indices.astype(index_dtype).tofile(indices_path)

        # --- SVD residual correction ---
        svd_bytes = 0
        cosine_with_svd = cosine_with_sidecar
        svd_rank_actual = 0

        if policy.svd_residual_rank > 0:
            # Build VQ+sidecar reconstruction as 2D
            if num_outliers > 0:
                # patched is the flat VQ+sidecar array (already computed above)
                recon_2d = patched.reshape(rows, cols)
            else:
                recon_2d = reconstructed.reshape(rows, cols)

            residual = tensor - recon_2d
            svd_rank_actual = min(policy.svd_residual_rank, min(rows, cols))
            U, s, Vt = np.linalg.svd(residual, full_matrices=False)
            U_r = U[:, :svd_rank_actual].astype(np.float32)
            s_r = s[:svd_rank_actual].astype(np.float32)
            Vt_r = Vt[:svd_rank_actual, :].astype(np.float32)

            np.save(out_dir / "svd_U.npy", U_r)
            np.save(out_dir / "svd_s.npy", s_r)
            np.save(out_dir / "svd_Vt.npy", Vt_r)

            svd_bytes = sum(
                (out_dir / f).stat().st_size
                for f in ("svd_U.npy", "svd_s.npy", "svd_Vt.npy")
            )

            corrected = recon_2d + (U_r * s_r[None, :]) @ Vt_r
            cosine_with_svd = float(_cosine(flat, corrected.ravel()))

        # --- Write meta.json ---
        index_dtype_str = "uint16" if index_dtype == np.uint16 else "uint8"
        meta = {
            "format_version": "cdna_v3",
            "tensor_name": tensor_name,
            "shape": [rows, cols],
            "dtype": "float32",
            "index_dtype": index_dtype_str,
            "storage_mode": policy.storage_mode,
            "tensor_class": policy.tensor_class.value,
            "n_clusters": policy.n_clusters,
            "percentile": policy.percentile,
            "use_kmeans": policy.use_kmeans,
            "sidecar_enabled": policy.sidecar_enabled,
            "block_rows": policy.block_rows,
            "vector_dim": vector_dim,
            "codebook_sha256": codebook_sha256,
            "svd_residual_rank": svd_rank_actual,
            "source_artifact": source_artifact,
            "created_utc": datetime.now(timezone.utc).isoformat(),
        }
        (out_dir / "meta.json").write_text(json.dumps(meta, indent=2))

        # --- Write stats.json ---
        indices_bytes = indices_path.stat().st_size
        codebook_bytes = codebook_path.stat().st_size
        compressed_bytes = indices_bytes + codebook_bytes + sidecar_bytes + svd_bytes

        stats = {
            "tensor_name": tensor_name,
            "shape": [rows, cols],
            "original_bytes": tensor.nbytes,
            "compressed_bytes": compressed_bytes,
            "compression_ratio": round(tensor.nbytes / max(1, compressed_bytes), 2),
            "cosine_no_sidecar": round(cosine_no_sidecar, 6),
            "cosine_with_sidecar": round(cosine_with_sidecar, 6),
            "cosine_with_svd": round(cosine_with_svd, 6),
            "max_abs_diff": round(max_abs_diff, 6),
            "num_outliers": num_outliers,
            "sidecar_bytes": sidecar_bytes,
            "svd_residual_rank": svd_rank_actual,
            "svd_bytes": svd_bytes,
        }
        (out_dir / "stats.json").write_text(json.dumps(stats, indent=2))

        return stats

    def _build_vector_codebook_and_assign(
        self, tensor: np.ndarray, n_clusters: int, vector_dim: int
    ) -> tuple[np.ndarray, np.ndarray]:
        """Build vector codebook and assign indices for multi-dim VQ.

        Args:
            tensor: [rows, cols] float32 tensor
            n_clusters: Number of codebook entries
            vector_dim: Dimensionality of each vector (2, 4, etc.)

        Returns:
            (codebook [k, d], indices [rows * cols // d])
        """
        from helix_substrate.cdna_encoder import _vector_kmeans

        rows, cols = tensor.shape
        # Group adjacent elements along in_features (row-wise, memory-contiguous)
        vectors = tensor.reshape(-1, vector_dim)  # [rows * cols/d, d]
        n_vectors = len(vectors)

        # Subsample for codebook fitting (same approach as scalar path)
        max_fit_samples = min(n_vectors, _KMEANS_MAX_SAMPLES)
        if n_vectors > max_fit_samples:
            rng = np.random.RandomState(42)
            fit_idx = rng.choice(n_vectors, max_fit_samples, replace=False)
            fit_vectors = vectors[fit_idx]
        else:
            fit_vectors = vectors

        # Fit codebook on subsample
        codebook, _ = _vector_kmeans(fit_vectors, n_clusters, max_iters=10)

        # Assign ALL vectors using chunked matmul distance
        indices = self._chunked_vector_assign(vectors, codebook, n_clusters)
        return codebook.astype(np.float32), indices

    def _chunked_vector_assign(
        self, vectors: np.ndarray, codebook: np.ndarray, n_clusters: int,
        chunk_size: int = 50000,
    ) -> np.ndarray:
        """Assign vector codebook indices in chunks to avoid OOM.

        Uses GPU torch.cdist when available (150x faster), falls back to
        CPU matmul trick: ||x-c||^2 = ||x||^2 + ||c||^2 - 2*x@c.T.
        """
        index_dtype = np.uint8 if n_clusters <= 256 else np.uint16
        n = len(vectors)

        # GPU path: torch.cdist in chunks (500K vectors per chunk fits ~8GB VRAM)
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                gpu_chunk = 500000
                indices = np.empty(n, dtype=index_dtype)
                cb_g = torch.from_numpy(codebook.astype(np.float32)).cuda()
                for start in range(0, n, gpu_chunk):
                    end = min(start + gpu_chunk, n)
                    data_g = torch.from_numpy(
                        vectors[start:end].astype(np.float32)
                    ).cuda()
                    dists = torch.cdist(data_g, cb_g)
                    idx = torch.argmin(dists, dim=1).cpu().numpy().astype(index_dtype)
                    indices[start:end] = idx
                    del data_g, dists, idx
                del cb_g
                torch.cuda.empty_cache()
                return indices
        except (ImportError, torch.cuda.OutOfMemoryError):
            pass

        # CPU fallback
        indices = np.empty(n, dtype=index_dtype)
        C_sq = np.sum(codebook ** 2, axis=1, keepdims=True).T  # [1, k]

        for start in range(0, n, chunk_size):
            end = min(start + chunk_size, n)
            X_chunk = vectors[start:end].astype(np.float32)
            X_sq = np.sum(X_chunk ** 2, axis=1, keepdims=True)  # [chunk, 1]
            dists = X_sq + C_sq - 2.0 * (X_chunk @ codebook.T)  # [chunk, k]
            indices[start:end] = np.argmin(dists, axis=1).astype(index_dtype)

        return indices

    def _build_kmeans_codebook(
        self, flat: np.ndarray, n_clusters: int
    ) -> np.ndarray:
        """Build codebook using k-means on subsampled data."""
        from helix_substrate.cdna_encoder import _simple_kmeans

        # Subsample to avoid OOM in k-means
        if len(flat) > _KMEANS_MAX_SAMPLES:
            rng = np.random.RandomState(42)
            sample_idx = rng.choice(len(flat), _KMEANS_MAX_SAMPLES, replace=False)
            sample = flat[sample_idx]
        else:
            sample = flat

        codebook, _ = _simple_kmeans(sample, n_clusters, max_iters=10)
        return codebook.astype(np.float32)

    def _build_uniform_codebook(
        self, flat: np.ndarray, n_clusters: int
    ) -> np.ndarray:
        """Build uniform linear codebook."""
        vmin, vmax = float(flat.min()), float(flat.max())
        return np.linspace(vmin, vmax, n_clusters).astype(np.float32)

    def _chunked_assign(
        self, flat: np.ndarray, codebook: np.ndarray,
        index_dtype: np.dtype = np.uint8,
    ) -> np.ndarray:
        """Assign codebook indices in chunks to avoid OOM.

        Uses GPU if available for massive speedup on large tensors.
        """
        try:
            import torch
            if torch.cuda.is_available():
                return self._gpu_assign(flat, codebook, index_dtype)
        except ImportError:
            pass

        indices = np.empty(len(flat), dtype=index_dtype)
        for start in range(0, len(flat), _ASSIGN_CHUNK_SIZE):
            end = min(start + _ASSIGN_CHUNK_SIZE, len(flat))
            chunk = flat[start:end]
            dists = np.abs(chunk[:, np.newaxis] - codebook)
            indices[start:end] = np.argmin(dists, axis=1).astype(index_dtype)
        return indices

    def _gpu_assign(
        self, flat: np.ndarray, codebook: np.ndarray,
        index_dtype: np.dtype = np.uint8,
    ) -> np.ndarray:
        """GPU-accelerated codebook assignment. ~100x faster than CPU."""
        import torch
        cb_t = torch.from_numpy(codebook).cuda()  # [k]
        indices = np.empty(len(flat), dtype=index_dtype)
        # Scale chunk to fit GPU memory: distance matrix = chunk * k * 4 bytes
        # Target ~512 MB peak for distance matrix (safe on 4GB T2000)
        chunk_size = max(1024, (512 * 1024 * 1024) // (len(codebook) * 4))
        for start in range(0, len(flat), chunk_size):
            end = min(start + chunk_size, len(flat))
            chunk_t = torch.from_numpy(flat[start:end]).cuda()
            dists = torch.abs(chunk_t.unsqueeze(1) - cb_t.unsqueeze(0))
            indices[start:end] = torch.argmin(dists, dim=1).cpu().numpy().astype(index_dtype)
        return indices




def _safe_name(tensor_name: str) -> str:
    """Convert tensor name to filesystem-safe string."""
    return tensor_name.replace("/", "_").replace(".", "_")


def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity between two flat arrays."""
    dot = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)
