"""
Weight Transforms — Function-preserving mathematical operations on weight matrices.

All transforms operate on pairs of adjacent weight matrices and preserve
the model's input-output behavior exactly (or within float precision).

Weight Transforms:
- permute_neurons: Random neuron reordering (mathematically exact)
- rotate_basis: Random orthogonal basis rotation (~1e-7 float error)
- scale_weights: Per-neuron random scaling with inverse compensation (exact)

Full-File Transforms (make ALL output files untraceable):
- randomize_metadata: Strip provenance from OsmoManifest
- obfuscate_knowledge: Noise-inject knowledge files (structural, functional, activations)
- obfuscate_architecture: Randomize layer names in ArchitectureSpec
- obfuscate_parameter_names: Rename all parameter keys to opaque hashes
"""

from __future__ import annotations

import copy
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

import hashlib

from model_osmosis.format.osmo_format import (
    ArchitectureSpec,
    LayerSpec,
    OsmoManifest,
)
from model_osmosis.obfuscation.graph import TransformGroup


def permute_neurons(
    params: Dict[str, np.ndarray],
    group: TransformGroup,
    rng: np.random.Generator,
) -> None:
    """Randomly reorder neurons in a hidden layer.

    For a hidden dimension of size N shared between W_current and W_next:
    - Generate a random permutation P of [0, 1, ..., N-1]
    - Reorder rows of W_current along axis_current
    - Reorder the bias (if any) with the same permutation
    - Reorder columns of W_next along axis_next

    This is mathematically exact — zero numerical error.
    Operates in-place on the params dict.
    """
    dim = group.dim_size
    perm = rng.permutation(dim)

    w_current = params[group.weight_key]
    w_next = params[group.next_weight_key]

    # Permute rows of current weight
    params[group.weight_key] = np.take(w_current, perm, axis=group.axis_current)

    # Permute parallel weights with the SAME permutation (e.g. up_proj in SwiGLU)
    for pk in group.parallel_weight_keys:
        if pk in params:
            params[pk] = np.take(params[pk], perm, axis=group.axis_current)

    # Permute bias if present
    if group.bias_key and group.bias_key in params:
        params[group.bias_key] = params[group.bias_key][perm]

    # Permute columns of next weight (inverse = same permutation on that axis)
    params[group.next_weight_key] = np.take(w_next, perm, axis=group.axis_next)

    # Permute normalization parameters
    for nk in group.norm_keys:
        if nk in params:
            params[nk] = params[nk][perm]


def rotate_basis(
    params: Dict[str, np.ndarray],
    group: TransformGroup,
    rng: np.random.Generator,
) -> None:
    """Apply a random orthogonal basis rotation to the hidden dimension.

    For adjacent layers with shared hidden dimension H:
    - Generate a random orthogonal matrix Q of shape (H, H)
    - W_current_new = Q @ W_current (along axis_current)
    - W_next_new = W_next @ Q^T (along axis_next)
    - bias_new = Q @ bias (if present)

    The computation Q^T @ Q = I ensures (W_next @ Q^T)(Q @ W_current) = W_next @ W_current.

    Introduces ~1e-7 float32 error. Entire computation done in float64
    internally for maximum precision; only casts back at the very end.
    Operates in-place on the params dict.
    """
    dim = group.dim_size

    # Generate random orthogonal matrix via QR decomposition
    random_matrix = rng.standard_normal((dim, dim))
    q_matrix, _ = np.linalg.qr(random_matrix)

    # Ensure proper rotation (det = +1), not reflection
    if np.linalg.det(q_matrix) < 0:
        q_matrix[:, 0] *= -1

    q_matrix = q_matrix.astype(np.float64)
    qt_matrix = q_matrix.T

    w_current = params[group.weight_key]
    w_next = params[group.next_weight_key]

    original_dtype_current = w_current.dtype
    original_dtype_next = w_next.dtype

    # All computations in float64 — only cast at the very end
    w_current_64 = w_current.astype(np.float64)
    w_next_64 = w_next.astype(np.float64)

    # Apply Q @ W_current along the appropriate axis
    if group.axis_current == 0:
        new_w_current = q_matrix @ w_current_64
        # Rotate parallel weights with the SAME Q (e.g. up_proj in SwiGLU)
        for pk in group.parallel_weight_keys:
            if pk in params:
                pk_64 = params[pk].astype(np.float64)
                params[pk] = (q_matrix @ pk_64).astype(params[pk].dtype)
    else:
        new_w_current = w_current_64 @ qt_matrix
        for pk in group.parallel_weight_keys:
            if pk in params:
                pk_64 = params[pk].astype(np.float64)
                params[pk] = (pk_64 @ qt_matrix).astype(params[pk].dtype)

    # Apply W_next @ Q^T along the appropriate axis
    if group.axis_next == 1:
        # W shape is [out, in], transform input dimension (cols)
        new_w_next = w_next_64 @ qt_matrix
    else:
        # W shape is [in, out], transform input dimension (rows)
        new_w_next = qt_matrix.T @ w_next_64

    # Cast back to original dtypes only at the very end
    params[group.weight_key] = new_w_current.astype(original_dtype_current)
    params[group.next_weight_key] = new_w_next.astype(original_dtype_next)

    # Transform bias: bias_new = Q @ bias
    if group.bias_key and group.bias_key in params:
        bias = params[group.bias_key]
        original_dtype_bias = bias.dtype
        bias_64 = bias.astype(np.float64)
        params[group.bias_key] = (q_matrix @ bias_64).astype(original_dtype_bias)

    # Transform norm parameters: norm_new = Q @ norm
    # Note: For LayerNorm, permutation is safe but rotation changes semantics.
    # We skip norm rotation here — norm params are only handled by permutation
    # and scaling, which are exact.


def scale_weights(
    params: Dict[str, np.ndarray],
    group: TransformGroup,
    rng: np.random.Generator,
    scale_range: Tuple[float, float] = (0.5, 2.0),
) -> None:
    """Apply per-neuron random scaling with inverse compensation.

    For each neuron i in the hidden dimension:
    - Sample a random scale factor s[i] from [scale_range[0], scale_range[1]]
    - Multiply row i of W_current by s[i]
    - Multiply bias[i] by s[i] (if present)
    - Divide column i of W_next by s[i]

    Since (W_next / s) @ (s * W_current) = W_next @ W_current, this is exact.
    Operates in-place on the params dict.
    """
    dim = group.dim_size
    low, high = scale_range
    scales = rng.uniform(low, high, size=dim).astype(np.float64)
    inv_scales = 1.0 / scales

    w_current = params[group.weight_key]
    w_next = params[group.next_weight_key]

    original_dtype_current = w_current.dtype
    original_dtype_next = w_next.dtype

    # Scale rows of W_current
    if group.axis_current == 0:
        scale_vec = scales.reshape(-1, 1).astype(original_dtype_current)
        params[group.weight_key] = w_current * scale_vec
        # Scale parallel weights with the SAME scales (e.g. up_proj in SwiGLU)
        for pk in group.parallel_weight_keys:
            if pk in params:
                params[pk] = params[pk] * scale_vec.astype(params[pk].dtype)
    else:
        scale_vec = scales.reshape(1, -1).astype(original_dtype_current)
        params[group.weight_key] = w_current * scale_vec
        for pk in group.parallel_weight_keys:
            if pk in params:
                params[pk] = params[pk] * scale_vec.astype(params[pk].dtype)

    # Scale bias
    if group.bias_key and group.bias_key in params:
        bias = params[group.bias_key]
        params[group.bias_key] = bias * scales.astype(bias.dtype)

    # Inverse-scale columns of W_next
    if group.axis_next == 1:
        # Shape [out, in] — inverse scale along cols (axis 1)
        inv_vec = inv_scales.reshape(1, -1).astype(original_dtype_next)
        params[group.next_weight_key] = w_next * inv_vec
    else:
        # Shape [in, out] — inverse scale along rows (axis 0)
        inv_vec = inv_scales.reshape(-1, 1).astype(original_dtype_next)
        params[group.next_weight_key] = w_next * inv_vec

    # Note: Normalization parameters (LayerNorm gamma/beta) are NOT scaled.
    # Scaling gamma without a corresponding inverse elsewhere would change
    # the model's output. Norm params are only safe to permute (reorder),
    # which is handled by permute_neurons.


def randomize_metadata(manifest: OsmoManifest) -> OsmoManifest:
    """Strip provenance information from the manifest.

    Clears:
    - original_repo
    - training_info (dataset hints, base model, hardware)
    - tags that reference source models
    - description (replaced with generic text)
    - created_at (regenerated)

    Returns a new OsmoManifest with cleaned metadata.
    """
    cleaned = copy.deepcopy(manifest)

    # Clear source references
    cleaned.original_repo = None
    cleaned.description = "Model processed with ModelOsmosis"

    # Regenerate timestamp
    cleaned.created_at = datetime.now(timezone.utc).isoformat()

    # Clear training info that could trace lineage
    cleaned.training_info.dataset_hints = []
    cleaned.training_info.base_model = None
    cleaned.training_info.hardware = None

    # Remove tags that reference specific models
    source_tag_patterns = [
        "merged", "slerp", "ties", "absorbed", "stacked",
        "llama", "mistral", "gpt", "qwen", "phi",
    ]
    cleaned.tags = [
        tag for tag in cleaned.tags
        if not any(p in tag.lower() for p in source_tag_patterns)
    ]

    return cleaned


def obfuscate_knowledge(
    knowledge: Dict[str, np.ndarray],
    rng: np.random.Generator,
    noise_scale: float = 0.01,
) -> Dict[str, np.ndarray]:
    """Obfuscate knowledge files so they cannot be used to trace the source model.

    Knowledge files contain:
    - structural weights (direct copies of source model weights)
    - functional I/O pairs (inputs/outputs that fingerprint model behavior)
    - intermediate activations (layer-by-layer fingerprints)

    When noise_scale > 0:
        Arrays are replaced with noise-injected versions that destroy
        traceability while keeping the same shapes and rough distributions.

    When noise_scale <= 0 (quality-preserving mode):
        Arrays are only permuted along axes (no noise added). This preserves
        exact values while preventing direct index-based matching.

    Args:
        knowledge: Dict of knowledge name → numpy array.
        rng: Random number generator.
        noise_scale: Scale of noise relative to data std (0.01 = 1%).
            Set to 0.0 for quality-preserving mode (permutation only).

    Returns:
        New dict with obfuscated knowledge arrays.
    """
    obfuscated: Dict[str, np.ndarray] = {}

    for key, arr in knowledge.items():
        if arr.size == 0:
            obfuscated[key] = arr.copy()
            continue

        result = arr.copy()

        # Randomly permute along the first axis (shuffles samples/layers)
        if arr.ndim >= 1 and arr.shape[0] > 1:
            perm = rng.permutation(arr.shape[0])
            result = result[perm]

        # Add noise only if noise_scale > 0 (non-quality-preserving mode)
        if noise_scale > 0:
            std = float(np.std(arr)) if arr.size > 1 else 1.0
            noise = rng.standard_normal(arr.shape).astype(arr.dtype) * (std * noise_scale)
            result = result + noise

        obfuscated[key] = result

    return obfuscated


def obfuscate_architecture(
    architecture: ArchitectureSpec,
    rng: np.random.Generator,
) -> ArchitectureSpec:
    """Randomize layer names in the architecture spec.

    Replaces identifiable layer names (e.g., 'layers.0.attention.q_proj.weight')
    with opaque hashes (e.g., 'L_a3f1b2c4') so the model family cannot be
    identified from the architecture file alone.

    Also strips any extra metadata fields that could leak provenance.

    Returns a new ArchitectureSpec with randomized names.
    """
    new_arch = copy.deepcopy(architecture)

    # Build a deterministic but random name mapping
    salt = rng.bytes(16).hex()

    new_layers: List[LayerSpec] = []
    for i, layer in enumerate(new_arch.layers):
        # Generate opaque name from index + random salt
        raw = f"{salt}_{i}_{layer.name}"
        hashed = hashlib.sha256(raw.encode()).hexdigest()[:8]
        new_name = f"L_{hashed}"

        new_layer = LayerSpec(
            name=new_name,
            layer_type=layer.layer_type,
            shape=layer.shape,
            dtype=layer.dtype,
            inputs=[],   # Clear connection graph
            outputs=[],
            attributes={},  # Clear extra attributes
        )
        new_layers.append(new_layer)

    new_arch.layers = new_layers

    # Clear extra metadata that could identify model family
    new_arch.extra = {}

    return new_arch


def obfuscate_auxiliary_files(
    auxiliary_files: Dict[str, Any],
    rng: np.random.Generator,
) -> Dict[str, Any]:
    """Obfuscate auxiliary files (tokenizer, config, preprocessor) to remove provenance.

    Strips or randomizes fields that could identify the source model:
    - config.json: clears _name_or_path, model repo references
    - tokenizer.json: randomizes added_tokens order, strips trainer info
    - tokenizer_config.json: clears name_or_path, auto_map references
    - special_tokens_map.json: kept (functional, but token content is generic)
    - generation_config.json: clears _from_model_config reference

    Args:
        auxiliary_files: Dict of filename -> content (JSON dicts or strings).
        rng: Random number generator.

    Returns:
        New dict with obfuscated auxiliary files.
    """
    obfuscated: Dict[str, Any] = {}

    for filename, data in auxiliary_files.items():
        if not isinstance(data, dict):
            # Non-dict files (raw text, bytes) — keep as-is
            obfuscated[filename] = data
            continue

        cleaned = copy.deepcopy(data)

        if filename == "config.json":
            cleaned = _obfuscate_config(cleaned, rng)
        elif filename == "tokenizer_config.json":
            cleaned = _obfuscate_tokenizer_config(cleaned, rng)
        elif filename == "tokenizer.json":
            cleaned = _obfuscate_tokenizer(cleaned, rng)
        elif filename == "generation_config.json":
            cleaned = _obfuscate_generation_config(cleaned, rng)
        elif filename == "special_tokens_map.json":
            # Functional file — keep as-is (tokens are needed for inference)
            pass
        else:
            # Unknown auxiliary file — strip any name_or_path references
            cleaned = _strip_provenance_fields(cleaned)

        obfuscated[filename] = cleaned

    return obfuscated


def _obfuscate_config(config: Dict[str, Any], rng: np.random.Generator) -> Dict[str, Any]:
    """Strip provenance from config.json."""
    # Remove fields that identify the source model
    provenance_keys = [
        "_name_or_path", "_commit_hash", "transformers_version",
        "auto_map", "_attn_implementation_autoset",
    ]
    for key in provenance_keys:
        config.pop(key, None)

    # Clear model_type only if obfuscating layer names too (it reveals architecture)
    # We keep it because it's needed for loading — but clear repo-specific fields

    return config


def _obfuscate_tokenizer_config(config: Dict[str, Any], rng: np.random.Generator) -> Dict[str, Any]:
    """Strip provenance from tokenizer_config.json."""
    provenance_keys = [
        "name_or_path", "auto_map", "_commit_hash",
    ]
    for key in provenance_keys:
        config.pop(key, None)

    return config


def _obfuscate_tokenizer(tokenizer: Dict[str, Any], rng: np.random.Generator) -> Dict[str, Any]:
    """Strip provenance from tokenizer.json (the full tokenizer spec)."""
    # Remove trainer-specific metadata
    tokenizer.pop("trainer", None)

    # Sort added_tokens by their explicit ID to normalize ordering
    # without breaking the ID mapping. Shuffling would break models
    # that rely on token position matching token ID.
    if "added_tokens" in tokenizer and isinstance(tokenizer["added_tokens"], list):
        tokens = tokenizer["added_tokens"]
        # Only sort if tokens have explicit 'id' fields (safe reordering)
        if tokens and isinstance(tokens[0], dict) and "id" in tokens[0]:
            tokenizer["added_tokens"] = sorted(tokens, key=lambda t: t["id"])

    return tokenizer


def _obfuscate_generation_config(config: Dict[str, Any], rng: np.random.Generator) -> Dict[str, Any]:
    """Strip provenance from generation_config.json."""
    provenance_keys = [
        "_from_model_config", "transformers_version", "_commit_hash",
    ]
    for key in provenance_keys:
        config.pop(key, None)

    return config


def _strip_provenance_fields(data: Dict[str, Any]) -> Dict[str, Any]:
    """Generic provenance stripping for unknown auxiliary files."""
    provenance_patterns = [
        "_name_or_path", "name_or_path", "_commit_hash",
        "auto_map", "transformers_version",
    ]
    for key in provenance_patterns:
        data.pop(key, None)
    return data


def obfuscate_parameter_names(
    params: Dict[str, np.ndarray],
    rng: np.random.Generator,
) -> Dict[str, np.ndarray]:
    """Rename all parameter keys to opaque hashes.

    Original names like 'model.layers.0.self_attn.q_proj.weight' reveal the
    model family. This replaces them with 'P_<hash>' while preserving the
    same weight arrays.

    The mapping is one-way — original names cannot be recovered.

    Args:
        params: Original parameter dict.
        rng: Random number generator.

    Returns:
        New dict with hashed parameter keys and same arrays.
    """
    salt = rng.bytes(32).hex()
    obfuscated: Dict[str, np.ndarray] = {}

    for key, arr in params.items():
        # Hash the salt + key + shape for a content-aware, non-sequential mapping.
        # Using shape prevents collisions when different params share similar names.
        shape_str = "x".join(str(d) for d in arr.shape)
        raw = f"{salt}:{key}:{shape_str}:{arr.dtype}"
        hashed = hashlib.sha256(raw.encode()).hexdigest()[:12]

        # Preserve weight/bias suffix for minimal compatibility
        if key.endswith(".weight") or key.endswith("_weight"):
            new_key = f"P_{hashed}.weight"
        elif key.endswith(".bias") or key.endswith("_bias"):
            new_key = f"P_{hashed}.bias"
        else:
            new_key = f"P_{hashed}"

        obfuscated[new_key] = arr

    return obfuscated
